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 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\n---\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---\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\n---\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",
        }
    ]