543 lines
19 KiB
Python
543 lines
19 KiB
Python
from __future__ import annotations
|
||
|
||
from types import SimpleNamespace
|
||
|
||
import pytest
|
||
import yuxi.agents.middlewares.subagent_task as subagent_task_middleware
|
||
from langchain_core.messages import AIMessage, HumanMessage
|
||
from langgraph.prebuilt.tool_node import ToolRuntime
|
||
from langgraph.types import Command
|
||
from yuxi.agents.buildin.chatbot.state import merge_subagent_runs
|
||
from yuxi.agents.middlewares.subagent_task import YuxiSubAgentMiddleware
|
||
from yuxi.utils.subagent_thread_utils import make_child_thread_id
|
||
from yuxi.repositories.agent_repository import SUB_AGENT_BACKEND_ID
|
||
|
||
|
||
class _ChildContext:
|
||
def __init__(self):
|
||
self.model = None
|
||
|
||
def update_from_dict(self, values: dict):
|
||
for key, value in values.items():
|
||
if hasattr(self, key):
|
||
setattr(self, key, value)
|
||
|
||
|
||
def _patch_subagent_run_tracking(monkeypatch):
|
||
async def create_run(_self, **kwargs):
|
||
return SimpleNamespace(
|
||
id=f"sub-run-{kwargs['tool_call_id']}",
|
||
request_id=f"sub-req-{kwargs['tool_call_id']}",
|
||
status="running",
|
||
parent_agent_run_id="parent-run",
|
||
created_at=None,
|
||
finished_at=None,
|
||
error_message=None,
|
||
), True
|
||
|
||
async def set_status(_self, run_id, status, *, error_type=None, error_message=None):
|
||
del error_type
|
||
return SimpleNamespace(
|
||
id=run_id,
|
||
request_id=run_id.replace("sub-run", "sub-req"),
|
||
status=status,
|
||
parent_agent_run_id="parent-run",
|
||
created_at=None,
|
||
finished_at=None,
|
||
error_message=error_message,
|
||
)
|
||
|
||
monkeypatch.setattr(YuxiSubAgentMiddleware, "_create_subagent_run", create_run)
|
||
monkeypatch.setattr(YuxiSubAgentMiddleware, "_set_subagent_run_status", set_status)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_create_task_middleware_loads_all_visible_subagents_when_empty(monkeypatch) -> None:
|
||
class _SessionContext:
|
||
async def __aenter__(self):
|
||
return object()
|
||
|
||
async def __aexit__(self, exc_type, exc, tb):
|
||
return None
|
||
|
||
class _UserRepository:
|
||
async def get_by_uid_with_db(self, _db, uid):
|
||
assert uid == "user-1"
|
||
return SimpleNamespace(uid="user-1", role="user")
|
||
|
||
class _AgentRepository:
|
||
def __init__(self, _db):
|
||
pass
|
||
|
||
async def list_visible_subagents(self, *, user):
|
||
assert user.uid == "user-1"
|
||
return [
|
||
SimpleNamespace(
|
||
slug="worker",
|
||
name="Worker",
|
||
description="work on scoped tasks",
|
||
backend_id=SUB_AGENT_BACKEND_ID,
|
||
config_json={},
|
||
),
|
||
SimpleNamespace(
|
||
slug="invalid",
|
||
name="Invalid",
|
||
description="invalid backend",
|
||
backend_id="ChatbotAgent",
|
||
config_json={},
|
||
),
|
||
]
|
||
|
||
async def get_visible_subagent_by_slug(self, *, slug, user):
|
||
raise AssertionError("empty subagents should load all visible subagents")
|
||
|
||
monkeypatch.setattr(
|
||
subagent_task_middleware,
|
||
"pg_manager",
|
||
SimpleNamespace(get_async_session_context=lambda: _SessionContext()),
|
||
)
|
||
monkeypatch.setattr(subagent_task_middleware, "UserRepository", _UserRepository)
|
||
monkeypatch.setattr(subagent_task_middleware, "AgentRepository", _AgentRepository)
|
||
|
||
middleware = await subagent_task_middleware.create_subagent_task_middleware(
|
||
SimpleNamespace(thread_id="parent-thread", uid="user-1", subagents=[])
|
||
)
|
||
|
||
assert isinstance(middleware, YuxiSubAgentMiddleware)
|
||
assert middleware.subagent_names == frozenset({"worker"})
|
||
assert middleware.transformers
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_task_tool_rejects_unconfigured_subagent() -> None:
|
||
middleware = YuxiSubAgentMiddleware(
|
||
parent_context=SimpleNamespace(thread_id="parent-thread", uid="user-1"),
|
||
subagents=[
|
||
SimpleNamespace(
|
||
slug="worker",
|
||
name="Worker",
|
||
description="work on scoped tasks",
|
||
backend_id=SUB_AGENT_BACKEND_ID,
|
||
config_json={},
|
||
)
|
||
],
|
||
)
|
||
runtime = ToolRuntime(
|
||
state={},
|
||
context=None,
|
||
tool_call_id="tool-1",
|
||
store=None,
|
||
stream_writer=lambda _: None,
|
||
config={},
|
||
)
|
||
|
||
result = await middleware.tools[0].ainvoke(
|
||
{"description": "do work", "subagent_type": "missing", "runtime": runtime}
|
||
)
|
||
|
||
assert result == "无法调用子智能体 missing，可用子智能体只有：`worker`"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_task_tool_invokes_subagent_with_child_scope(monkeypatch) -> None:
|
||
_patch_subagent_run_tracking(monkeypatch)
|
||
captured = {}
|
||
|
||
class _Graph:
|
||
async def ainvoke(self, state, *, config, context):
|
||
captured["state"] = state
|
||
captured["config"] = config
|
||
captured["context"] = context
|
||
return {
|
||
"messages": [AIMessage(content="child done")],
|
||
"artifacts": ["/home/gem/user-data/outputs/report.md"],
|
||
"todos": ["should not merge"],
|
||
}
|
||
|
||
class _Backend:
|
||
context_schema = _ChildContext
|
||
|
||
async def get_graph(self, *, context):
|
||
captured["graph_context"] = context
|
||
return _Graph()
|
||
|
||
monkeypatch.setattr(
|
||
subagent_task_middleware,
|
||
"_get_agent_backend",
|
||
lambda backend_id: _Backend() if backend_id == SUB_AGENT_BACKEND_ID else None,
|
||
)
|
||
times = iter(["2026-05-31T01:00:00Z", "2026-05-31T01:00:03Z"])
|
||
monkeypatch.setattr(subagent_task_middleware, "utc_isoformat", lambda: next(times))
|
||
|
||
middleware = YuxiSubAgentMiddleware(
|
||
parent_context=SimpleNamespace(
|
||
thread_id="child-runtime-thread",
|
||
parent_thread_id="parent-thread",
|
||
file_thread_id="parent-file-thread",
|
||
uid="user-1",
|
||
),
|
||
subagents=[
|
||
SimpleNamespace(
|
||
slug="worker.agent",
|
||
name="Worker",
|
||
description="work on scoped tasks",
|
||
backend_id=SUB_AGENT_BACKEND_ID,
|
||
config_json={"context": {"model": "provider:model", "subagents": ["nested"]}},
|
||
)
|
||
],
|
||
)
|
||
runtime = SimpleNamespace(
|
||
tool_call_id="tool-1",
|
||
state={
|
||
"messages": [HumanMessage(content="parent")],
|
||
"todos": ["parent todo"],
|
||
"activated_skills": ["parent-skill"],
|
||
"kept": "value",
|
||
},
|
||
config={
|
||
"callbacks": ["stream-callback"],
|
||
"tags": ["parent"],
|
||
"recursion_limit": 42,
|
||
"configurable": {"checkpoint_ns": "parent-ns", "__pregel_task_id": "parent-task"},
|
||
},
|
||
)
|
||
|
||
result = await middleware.tools[0].coroutine(
|
||
description="write a report",
|
||
subagent_type="worker.agent",
|
||
runtime=runtime,
|
||
)
|
||
|
||
child_thread_id = make_child_thread_id("parent-thread", "worker.agent", "tool-1")
|
||
assert isinstance(result, Command)
|
||
assert result.update["messages"][0].content == f"子智能体线程 ID: {child_thread_id}\n\nchild done"
|
||
assert result.update["messages"][0].tool_call_id == "tool-1"
|
||
assert result.update["artifacts"] == ["/home/gem/user-data/outputs/report.md"]
|
||
assert result.update["subagent_runs"] == [
|
||
{
|
||
"id": "tool-1",
|
||
"subagent_type": "worker.agent",
|
||
"subagent_name": "Worker",
|
||
"child_thread_id": child_thread_id,
|
||
"description": "write a report",
|
||
"created_at": "2026-05-31T01:00:00Z",
|
||
"run_id": "sub-run-tool-1",
|
||
"parent_agent_run_id": "parent-run",
|
||
"status": "completed",
|
||
"completed_at": "2026-05-31T01:00:03Z",
|
||
"result_preview": "child done",
|
||
"error": None,
|
||
"artifacts": ["/home/gem/user-data/outputs/report.md"],
|
||
}
|
||
]
|
||
assert "kept" not in captured["state"]
|
||
assert captured["state"]["parent_thread_id"] == "parent-thread"
|
||
assert captured["state"]["file_thread_id"] == "parent-file-thread"
|
||
assert captured["state"]["skills_thread_id"] == child_thread_id
|
||
assert captured["state"]["messages"] == [HumanMessage(content="write a report")]
|
||
assert "todos" not in captured["state"]
|
||
assert "activated_skills" not in captured["state"]
|
||
assert captured["config"]["callbacks"] == ["stream-callback"]
|
||
assert captured["config"]["tags"] == ["parent"]
|
||
assert captured["config"]["recursion_limit"] == 42
|
||
assert captured["config"]["configurable"] == {
|
||
"thread_id": child_thread_id,
|
||
"uid": "user-1",
|
||
"parent_thread_id": "parent-thread",
|
||
"file_thread_id": "parent-file-thread",
|
||
"skills_thread_id": child_thread_id,
|
||
"subagent_type": "worker.agent",
|
||
"subagent_thread_id": child_thread_id,
|
||
"subagent_tool_call_id": "tool-1",
|
||
"run_id": "sub-run-tool-1",
|
||
"request_id": "sub-req-tool-1",
|
||
"ls_agent_type": "subagent",
|
||
}
|
||
assert captured["context"] is captured["graph_context"]
|
||
assert captured["context"].model == "provider:model"
|
||
assert captured["context"].thread_id == child_thread_id
|
||
assert captured["context"].parent_thread_id == "parent-thread"
|
||
assert captured["context"].file_thread_id == "parent-file-thread"
|
||
assert captured["context"].skills_thread_id == child_thread_id
|
||
assert not hasattr(captured["context"], "subagents")
|
||
assert captured["context"].is_subagent_runtime is True
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_task_tool_records_failed_subagent_run(monkeypatch) -> None:
|
||
_patch_subagent_run_tracking(monkeypatch)
|
||
|
||
class _Graph:
|
||
async def ainvoke(self, state, *, config, context):
|
||
del state, config, context
|
||
raise RuntimeError("child boom")
|
||
|
||
class _Backend:
|
||
context_schema = _ChildContext
|
||
|
||
async def get_graph(self, *, context):
|
||
del context
|
||
return _Graph()
|
||
|
||
monkeypatch.setattr(
|
||
subagent_task_middleware,
|
||
"_get_agent_backend",
|
||
lambda backend_id: _Backend() if backend_id == SUB_AGENT_BACKEND_ID else None,
|
||
)
|
||
times = iter(["2026-05-31T02:00:00Z", "2026-05-31T02:00:04Z"])
|
||
monkeypatch.setattr(subagent_task_middleware, "utc_isoformat", lambda: next(times))
|
||
|
||
middleware = YuxiSubAgentMiddleware(
|
||
parent_context=SimpleNamespace(thread_id="parent-thread", uid="user-1"),
|
||
subagents=[
|
||
SimpleNamespace(
|
||
slug="worker",
|
||
name="Worker",
|
||
description="work on scoped tasks",
|
||
backend_id=SUB_AGENT_BACKEND_ID,
|
||
config_json={},
|
||
)
|
||
],
|
||
)
|
||
runtime = SimpleNamespace(tool_call_id="tool-1", state={}, config={})
|
||
|
||
result = await middleware.tools[0].coroutine(
|
||
description="write a report",
|
||
subagent_type="worker",
|
||
runtime=runtime,
|
||
)
|
||
|
||
assert isinstance(result, Command)
|
||
child_thread_id = make_child_thread_id("parent-thread", "worker", "tool-1")
|
||
assert (
|
||
result.update["messages"][0].content
|
||
== f"子智能体线程 ID: {child_thread_id}\n\n子智能体 worker 调用失败：child boom"
|
||
)
|
||
assert result.update["subagent_runs"] == [
|
||
{
|
||
"id": "tool-1",
|
||
"subagent_type": "worker",
|
||
"subagent_name": "Worker",
|
||
"child_thread_id": child_thread_id,
|
||
"description": "write a report",
|
||
"created_at": "2026-05-31T02:00:00Z",
|
||
"run_id": "sub-run-tool-1",
|
||
"parent_agent_run_id": "parent-run",
|
||
"status": "failed",
|
||
"completed_at": "2026-05-31T02:00:04Z",
|
||
"result_preview": "子智能体 worker 调用失败：child boom",
|
||
"error": "child boom",
|
||
"artifacts": [],
|
||
}
|
||
]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_task_tool_continues_existing_subagent_thread(monkeypatch) -> None:
|
||
_patch_subagent_run_tracking(monkeypatch)
|
||
captured = {}
|
||
|
||
class _Graph:
|
||
async def ainvoke(self, state, *, config, context):
|
||
captured["state"] = state
|
||
captured["config"] = config
|
||
captured["context"] = context
|
||
return {"messages": [AIMessage(content="continued done")]}
|
||
|
||
class _Backend:
|
||
context_schema = _ChildContext
|
||
|
||
async def get_graph(self, *, context):
|
||
captured["graph_context"] = context
|
||
return _Graph()
|
||
|
||
monkeypatch.setattr(
|
||
subagent_task_middleware,
|
||
"_get_agent_backend",
|
||
lambda backend_id: _Backend() if backend_id == SUB_AGENT_BACKEND_ID else None,
|
||
)
|
||
times = iter(["2026-05-31T03:00:00Z", "2026-05-31T03:00:05Z"])
|
||
monkeypatch.setattr(subagent_task_middleware, "utc_isoformat", lambda: next(times))
|
||
|
||
child_thread_id = make_child_thread_id("parent-thread", "worker.agent", "tool-old")
|
||
middleware = YuxiSubAgentMiddleware(
|
||
parent_context=SimpleNamespace(thread_id="parent-thread", uid="user-1"),
|
||
subagents=[
|
||
SimpleNamespace(
|
||
slug="worker.agent",
|
||
name="Worker",
|
||
description="work on scoped tasks",
|
||
backend_id=SUB_AGENT_BACKEND_ID,
|
||
config_json={},
|
||
)
|
||
],
|
||
)
|
||
runtime = SimpleNamespace(
|
||
tool_call_id="tool-2",
|
||
state={
|
||
"subagent_runs": [
|
||
{
|
||
"id": "tool-old",
|
||
"subagent_type": "worker.agent",
|
||
"child_thread_id": child_thread_id,
|
||
"status": "completed",
|
||
}
|
||
],
|
||
"kept": "parent-value",
|
||
},
|
||
config={"configurable": {"checkpoint_ns": "parent-ns"}},
|
||
)
|
||
|
||
result = await middleware.tools[0].coroutine(
|
||
description="continue the report",
|
||
subagent_type="worker.agent",
|
||
runtime=runtime,
|
||
thread_id=child_thread_id,
|
||
)
|
||
|
||
assert isinstance(result, Command)
|
||
assert result.update["messages"][0].content == f"子智能体线程 ID: {child_thread_id}\n\ncontinued done"
|
||
assert result.update["subagent_runs"] == [
|
||
{
|
||
"id": "tool-2",
|
||
"subagent_type": "worker.agent",
|
||
"subagent_name": "Worker",
|
||
"child_thread_id": child_thread_id,
|
||
"description": "continue the report",
|
||
"created_at": "2026-05-31T03:00:00Z",
|
||
"run_id": "sub-run-tool-2",
|
||
"parent_agent_run_id": "parent-run",
|
||
"status": "completed",
|
||
"completed_at": "2026-05-31T03:00:05Z",
|
||
"result_preview": "continued done",
|
||
"error": None,
|
||
"artifacts": [],
|
||
}
|
||
]
|
||
assert captured["state"] == {
|
||
"parent_thread_id": "parent-thread",
|
||
"file_thread_id": "parent-thread",
|
||
"skills_thread_id": child_thread_id,
|
||
"messages": [HumanMessage(content="continue the report")],
|
||
}
|
||
assert captured["config"]["configurable"] == {
|
||
"thread_id": child_thread_id,
|
||
"uid": "user-1",
|
||
"parent_thread_id": "parent-thread",
|
||
"file_thread_id": "parent-thread",
|
||
"skills_thread_id": child_thread_id,
|
||
"subagent_type": "worker.agent",
|
||
"subagent_thread_id": child_thread_id,
|
||
"subagent_tool_call_id": "tool-2",
|
||
"run_id": "sub-run-tool-2",
|
||
"request_id": "sub-req-tool-2",
|
||
"ls_agent_type": "subagent",
|
||
}
|
||
assert captured["context"] is captured["graph_context"]
|
||
assert captured["context"].thread_id == child_thread_id
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_task_tool_rejects_invalid_continuation_thread(monkeypatch) -> None:
|
||
graph_called = False
|
||
|
||
class _Backend:
|
||
context_schema = _ChildContext
|
||
|
||
async def get_graph(self, *, context):
|
||
nonlocal graph_called
|
||
del context
|
||
graph_called = True
|
||
raise AssertionError("invalid continuation should not invoke graph")
|
||
|
||
async def reject_continuation(_self, **kwargs):
|
||
raise ValueError(f"无法继续子智能体线程 {kwargs['child_thread_id']}：当前对话中没有找到对应的子智能体运行记录")
|
||
|
||
monkeypatch.setattr(
|
||
subagent_task_middleware,
|
||
"_get_agent_backend",
|
||
lambda backend_id: _Backend() if backend_id == SUB_AGENT_BACKEND_ID else None,
|
||
)
|
||
monkeypatch.setattr(YuxiSubAgentMiddleware, "_create_subagent_run", reject_continuation)
|
||
|
||
middleware = YuxiSubAgentMiddleware(
|
||
parent_context=SimpleNamespace(thread_id="parent-thread", uid="user-1", run_id="parent-run"),
|
||
subagents=[
|
||
SimpleNamespace(
|
||
slug="worker",
|
||
name="Worker",
|
||
description="work on scoped tasks",
|
||
backend_id=SUB_AGENT_BACKEND_ID,
|
||
config_json={},
|
||
)
|
||
],
|
||
)
|
||
|
||
unknown_thread_id = "opaque-child-thread"
|
||
runtime = SimpleNamespace(tool_call_id="tool-2", state={}, config={})
|
||
result = await middleware.tools[0].coroutine(
|
||
description="continue",
|
||
subagent_type="worker",
|
||
runtime=runtime,
|
||
thread_id=unknown_thread_id,
|
||
)
|
||
|
||
assert result == f"无法继续子智能体线程 {unknown_thread_id}：当前对话中没有找到对应的子智能体运行记录"
|
||
assert graph_called is False
|
||
|
||
|
||
def test_make_child_thread_id_fits_agent_run_thread_column() -> None:
|
||
child_thread_id = make_child_thread_id(
|
||
"fa62c751-d124-476f-a58c-855890aebcc4",
|
||
"agent-with-a-very-long-slug-that-would-overflow-the-column",
|
||
"019e86570b418b4ea6b5aee3ef87b1fa",
|
||
)
|
||
|
||
assert len(child_thread_id) <= 64
|
||
assert child_thread_id == make_child_thread_id(
|
||
"fa62c751-d124-476f-a58c-855890aebcc4",
|
||
"agent-with-a-very-long-slug-that-would-overflow-the-column",
|
||
"019e86570b418b4ea6b5aee3ef87b1fa",
|
||
)
|
||
|
||
|
||
def test_merge_subagent_runs_reuses_child_thread_entry() -> None:
|
||
child_thread_id = make_child_thread_id("parent-thread", "worker", "tool-old")
|
||
|
||
merged = merge_subagent_runs(
|
||
[
|
||
{
|
||
"id": "tool-old",
|
||
"subagent_type": "worker",
|
||
"subagent_name": "Worker",
|
||
"child_thread_id": child_thread_id,
|
||
"description": "first task",
|
||
"status": "completed",
|
||
"created_at": "2026-05-31T01:00:00Z",
|
||
}
|
||
],
|
||
[
|
||
{
|
||
"id": "tool-new",
|
||
"subagent_type": "worker",
|
||
"subagent_name": "Worker",
|
||
"child_thread_id": child_thread_id,
|
||
"description": "continue task",
|
||
"status": "completed",
|
||
"created_at": "2026-05-31T02:00:00Z",
|
||
}
|
||
],
|
||
)
|
||
|
||
assert merged == [
|
||
{
|
||
"id": "tool-new",
|
||
"subagent_type": "worker",
|
||
"subagent_name": "Worker",
|
||
"child_thread_id": child_thread_id,
|
||
"description": "continue task",
|
||
"status": "completed",
|
||
"created_at": "2026-05-31T02:00:00Z",
|
||
}
|
||
]
|