290 lines
10 KiB
Python
290 lines
10 KiB
Python
from __future__ import annotations
|
||
|
||
from types import SimpleNamespace
|
||
|
||
import pytest
|
||
import yuxi.agents.middlewares.subagent_task_middleware 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.middlewares.subagent_task_middleware import YuxiSubAgentMiddleware
|
||
from yuxi.agents.subagent_thread 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)
|
||
|
||
|
||
@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:
|
||
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 == "child 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",
|
||
"status": "completed",
|
||
"completed_at": "2026-05-31T01:00:03Z",
|
||
"result_preview": "child done",
|
||
"error": None,
|
||
"artifacts": ["/home/gem/user-data/outputs/report.md"],
|
||
}
|
||
]
|
||
assert captured["state"]["kept"] == "value"
|
||
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,
|
||
"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:
|
||
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)
|
||
assert result.update["messages"][0].content == "子智能体 worker 调用失败：child boom"
|
||
assert result.update["subagent_runs"] == [
|
||
{
|
||
"id": "tool-1",
|
||
"subagent_type": "worker",
|
||
"subagent_name": "Worker",
|
||
"child_thread_id": make_child_thread_id(
|
||
"parent-thread", "worker", "tool-1"
|
||
),
|
||
"description": "write a report",
|
||
"created_at": "2026-05-31T02:00:00Z",
|
||
"status": "failed",
|
||
"completed_at": "2026-05-31T02:00:04Z",
|
||
"result_preview": "子智能体 worker 调用失败：child boom",
|
||
"error": "child boom",
|
||
"artifacts": [],
|
||
}
|
||
]
|