PlantAI-dev/tests/unit/test_router_analysis_handoff_boundary.py

from __future__ import annotations

import asyncio
import json
from types import SimpleNamespace
from typing import Any

import pytest
from pydantic import BaseModel

from gfa_core.mas.agents.conversation_tools import ConversationToolRegistry
from gfa_core.mas.agents.router_runtime import (
    RouteDecision,
    RouterRuntime,
)
from gfa_core.mas.llm import LLMResponse, LLMToolCall
from gfa_core.mas.workflow_models import GoalLedger


class _PassThroughTools:
    @staticmethod
    def resolve_analysis_handoff_request(
        query: str,
        *,
        workflow_id: str | None = None,
        guidance: bool = False,
        language: str = "zh",
    ) -> str:
        return query


class _ContextEnrichingTools:
    @staticmethod
    def resolve_analysis_handoff_request(
        query: str,
        *,
        workflow_id: str | None = None,
        guidance: bool = False,
        language: str = "zh",
    ) -> str:
        return (
            f"{query}\n已确认上下文：当前 job 的 subfamily_confirm 已成功完成，"
            "家族成员和亚家族分类表已确认。"
        )


class _IntentArgs(BaseModel):
    query: str


class _SequenceLLM:
    def __init__(self, responses: list[LLMResponse]) -> None:
        self.responses = list(responses)

    async def acall(self, *_args: Any, **_kwargs: Any) -> LLMResponse:
        return self.responses.pop(0)


def _response(content: str = "", *, tool_calls: list[LLMToolCall] | None = None) -> LLMResponse:
    return LLMResponse(
        content=content,
        model="test-model",
        provider="test-provider",
        usage={},
        tool_calls=list(tool_calls or []),
    )


@pytest.mark.parametrize(
    ("query", "workflow_id", "clarification_question"),
    (
        (
            "执行rnaseq",
            "rnaseq_whole_genome",
            "请选择全基因组还是基因家族成员RNA-seq",
        ),
        ("执行全基因组rnaseq", "rnaseq_whole_genome", None),
        ("对已确认家族成员执行rnaseq", "rnaseq_family_member", None),
        ("执行motif分析", "gene_family", None),
    ),
)
def test_new_analysis_proposal_preserves_original_request_without_main_agent_planning(
    monkeypatch: pytest.MonkeyPatch,
    query: str,
    workflow_id: str,
    clarification_question: str | None,
) -> None:
    runtime = RouterRuntime(object(), _PassThroughTools(), max_iterations=5)  # type: ignore[arg-type]

    async def decide_route(**_: Any) -> tuple[RouteDecision, int, list[dict[str, Any]]]:
        return (
            RouteDecision(
                goal="new_analysis",
                needs=["analysis_intent"],
                subject=query,
                clarification_question=clarification_question,
                workflow_id=workflow_id,  # type: ignore[arg-type]
                analysis_request_kind="execution",
            ),
            1,
            [],
        )

    monkeypatch.setattr(runtime, "_decide_route", decide_route)

    outcome = asyncio.run(
        runtime.run(
            query=query,
            history=[],
            session_id="analysis-handoff-session",
            job_id="job-analysis-handoff",
            goal_ledger=GoalLedger(
                conversation_id="analysis-handoff-session",
                revision=0,
                user_terms={"active_job_id": "job-analysis-handoff"},
            ),
        )
    )

    assert outcome.action == "clarify"
    assert outcome.proposal is not None
    assert outcome.proposal["request"] == query
    assert outcome.proposal["original_request"] == query
    assert "workflow_id" not in outcome.proposal
    assert "全基因组还是" not in outcome.message


def test_context_enrichment_does_not_replace_original_analysis_handoff() -> None:
    runtime = RouterRuntime(object(), _ContextEnrichingTools(), max_iterations=5)  # type: ignore[arg-type]

    async def decide_route(**_: Any) -> tuple[RouteDecision, int, list[dict[str, Any]]]:
        return (
            RouteDecision(
                goal="new_analysis",
                needs=["analysis_intent"],
                subject="执行kaks分析",
                workflow_id="gene_family",
                analysis_request_kind="execution",
            ),
            1,
            [],
        )

    runtime._decide_route = decide_route  # type: ignore[method-assign]
    outcome = asyncio.run(
        runtime.run(
            query="执行kaks分析",
            history=[],
            session_id="kaks-handoff-session",
            job_id="job-kaks-handoff",
            goal_ledger=GoalLedger(
                conversation_id="kaks-handoff-session",
                revision=0,
                user_terms={"active_job_id": "job-kaks-handoff"},
            ),
        )
    )

    assert outcome.proposal is not None
    assert "subfamily_confirm" in outcome.proposal["request"]
    assert outcome.proposal["original_request"] == "执行kaks分析"


def test_confirmed_analysis_proposal_skips_second_confirmation() -> None:
    job_id = "job-confirmed-analysis-handoff"
    enriched_request = (
        "执行Ft的PPI和kaks分析\n"
        "已确认上下文：基因家族 SPL；目标物种 Ft；参考物种 At。"
    )
    recognized_queries: list[str] = []
    registry = ConversationToolRegistry()
    registry.register(
        name="recognize_analysis_intent",
        description="Recognize an analysis intent.",
        args_model=_IntentArgs,
        handler=lambda args: recognized_queries.append(args.query)
        or {
            "gene_family": "SPL",
            "stages": ["ppi", "kaks"],
            "primary_stage": "ppi",
            "params": {"targetSpecies": ["Ft"], "referenceSpecies": ["At"]},
        },
        permission="pure",
    )
    llm = _SequenceLLM(
        [
            _response(
                tool_calls=[
                    LLMToolCall(
                        id="recognize-confirmed-analysis-1",
                        name="recognize_analysis_intent",
                        arguments=json.dumps(
                            {"query": enriched_request}, ensure_ascii=False
                        ),
                    )
                ]
            ),
            _response(
                '{"status":"ok","action":"propose_analysis",'
                '"message":"已交给现有分析系统。"}'
            ),
        ]
    )
    ledger = GoalLedger(
        conversation_id="session-confirmed-analysis-handoff",
        revision=1,
        user_terms={
            "active_job_id": job_id,
            "pending_proposal": {
                "proposal_id": "proposal-confirmed-analysis",
                "request": enriched_request,
                "original_request": "执行Ft的PPI和kaks分析",
                "subject": "执行Ft的PPI和kaks分析",
                "job_id": job_id,
            },
        },
    )

    outcome = asyncio.run(
        RouterRuntime(llm, registry, max_iterations=5).run(
            query="继续",
            history=[],
            session_id="session-confirmed-analysis-handoff",
            job_id=job_id,
            goal_ledger=ledger,
        )
    )

    assert outcome.action == "propose_analysis"
    assert outcome.intent is not None
    assert outcome.intent["stages"] == ["ppi", "kaks"]
    assert recognized_queries == [enriched_request]


def test_new_analysis_decision_discards_main_agent_workflow_and_clarification() -> None:
    decision = RouteDecision(
        goal="new_analysis",
        needs=["analysis_intent"],
        subject="执行rnaseq",
        clarification_question="请选择RNA-seq模式",
        workflow_id="rnaseq_whole_genome",
        analysis_request_kind="execution",
    )

    normalized = RouterRuntime._analysis_handoff_decision(decision)

    assert normalized.goal == "new_analysis"
    assert normalized.needs == ["analysis_intent"]
    assert normalized.subject == decision.subject
    assert normalized.workflow_id is None
    assert normalized.clarification_question is None
    assert normalized.analysis_request_kind == "execution"


def test_router_response_language_does_not_rewrite_business_query() -> None:
    query = "执行接下来的分析"
    messages = RouterRuntime._initial_messages(
        query=query,
        history=[],
        session_id="english-router-session",
        job_id="job-english-router",
        decision=RouteDecision(
            goal="general",
            needs=[],
            subject=query,
        ),
        goal_ledger=None,
        language="en",
    )

    assert messages[-1].content == query
    assert "Respond to the user in English." in messages[0].content


def test_router_generated_model_failure_is_not_replayed_as_history() -> None:
    failure_message = (
        "Router 调用 DeepSeek 失败（HTTP 402，模型 deepseek-v4-flash，"
        "已尝试 3 次，错误码 invalid_request_error）：Insufficient Balance"
    )
    history = [
        {"role": "user", "content": "hello"},
        {"role": "assistant", "content": failure_message},
        {"role": "assistant", "content": "正常历史回复"},
    ]
    decision = RouteDecision(goal="general", needs=[], subject="hello")

    route_messages = RouterRuntime._route_messages(
        query="hello",
        history=history,
        session_id="router-history-session",
        job_id="job-router-history",
        goal_ledger=None,
    )
    response_messages = RouterRuntime._initial_messages(
        query="hello",
        history=history,
        session_id="router-history-session",
        job_id="job-router-history",
        decision=decision,
        goal_ledger=None,
    )

    for messages in (route_messages, response_messages):
        contents = [message.content for message in messages]
        assert failure_message not in contents
        assert "正常历史回复" in contents


def test_route_decision_retries_transient_model_failures(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class TransientPaymentGatewayFailure(RuntimeError):
        status_code = 402

    class TransientLLM:
        calls = 0

        async def acall(self, *_: Any, **__: Any) -> Any:
            self.calls += 1
            if self.calls < 3:
                raise TransientPaymentGatewayFailure("temporary billing gateway failure")
            return SimpleNamespace(
                content=(
                    '{"goal":"new_analysis","needs":["analysis_intent"],'
                    '"subject":"执行数据准备","analysis_request_kind":"execution"}'
                )
            )

    async def no_sleep(_: float) -> None:
        return None

    llm = TransientLLM()
    runtime = RouterRuntime(llm, _PassThroughTools(), max_iterations=5)  # type: ignore[arg-type]
    monkeypatch.setattr(asyncio, "sleep", no_sleep)

    decision, calls, trace = asyncio.run(
        runtime._decide_route(
            query="执行数据准备",
            history=[],
            session_id="retry-session",
            job_id="job-retry",
            goal_ledger=None,
        )
    )

    assert llm.calls == 3
    assert calls == 1
    assert trace == []
    assert decision is not None
    assert decision.goal == "new_analysis"


def test_route_decision_does_not_retry_authentication_failure() -> None:
    class AuthenticationFailure(RuntimeError):
        status_code = 401

    class UnauthorizedLLM:
        calls = 0

        async def acall(self, *_: Any, **__: Any) -> Any:
            self.calls += 1
            raise AuthenticationFailure("unauthorized")

    llm = UnauthorizedLLM()
    runtime = RouterRuntime(llm, _PassThroughTools(), max_iterations=5)  # type: ignore[arg-type]

    decision, calls, trace = asyncio.run(
        runtime._decide_route(
            query="执行数据准备",
            history=[],
            session_id="auth-session",
            job_id="job-auth",
            goal_ledger=None,
        )
    )

    assert llm.calls == 1
    assert decision is None
    assert calls == 1
    assert trace == []


def test_route_failure_exposes_sanitized_upstream_reason(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class StructuredPaymentGatewayFailure(RuntimeError):
        status_code = 402
        request_id = "req-router-402"
        body = {
            "error": {
                "code": "request_policy_rejected",
                "message": (
                    "the concrete upstream reason; "
                    "Authorization: Bearer must-not-leak"
                ),
                "api_key": "must-not-leak",
            }
        }

    class FailingLLM:
        _default_model = "deepseek-v4-flash"
        calls = 0

        async def acall(self, *_: Any, **__: Any) -> Any:
            self.calls += 1
            raise StructuredPaymentGatewayFailure("generic SDK error")

    async def no_sleep(_: float) -> None:
        return None

    llm = FailingLLM()
    runtime = RouterRuntime(llm, _PassThroughTools(), max_iterations=5)  # type: ignore[arg-type]
    monkeypatch.setattr(asyncio, "sleep", no_sleep)

    outcome = asyncio.run(
        runtime.run(
            query="执行数据准备",
            history=[],
            session_id="failure-detail-session",
            job_id="job-failure-detail",
            goal_ledger=None,
        )
    )

    assert llm.calls == 3
    assert outcome.status == "agent_unavailable"
    assert "HTTP 402" in outcome.message
    assert "request_policy_rejected" in outcome.message
    assert "the concrete upstream reason" in outcome.message
    assert "deepseek-v4-flash" in outcome.message
    assert "3 次" in outcome.message
    assert "req-router-402" in outcome.message
    assert "must-not-leak" not in outcome.message


def test_route_failure_reads_flat_openai_error_body(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class FlatOpenAIError(RuntimeError):
        status_code = 402
        body = {
            "code": "invalid_request_error",
            "message": "Insufficient Balance",
            "type": "unknown_error",
        }

    class FailingLLM:
        _default_model = "deepseek-v4-flash"

        async def acall(self, *_: Any, **__: Any) -> Any:
            raise FlatOpenAIError("serialized SDK fallback")

    async def no_sleep(_: float) -> None:
        return None

    runtime = RouterRuntime(FailingLLM(), _PassThroughTools(), max_iterations=5)  # type: ignore[arg-type]
    monkeypatch.setattr(asyncio, "sleep", no_sleep)

    outcome = asyncio.run(
        runtime.run(
            query="执行数据准备",
            history=[],
            session_id="flat-openai-error-session",
            job_id="job-flat-openai-error",
            goal_ledger=None,
        )
    )

    assert outcome.status == "agent_unavailable"
    assert "错误码 invalid_request_error" in outcome.message
    assert outcome.message.endswith("Insufficient Balance")
    assert "serialized SDK fallback" not in outcome.message