ForcePilot/backend/test/unit/middlewares/test_subagent_task_middleware.py

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": [],
        }
    ]