meida_backend/tests/unit/test_gateway.py

"""LLM gateway: routing, failover, circuit breaking, structured output, cost.

No test here touches the network — providers are fakes.
"""

from __future__ import annotations

import itertools
from collections.abc import Mapping, Sequence
from decimal import Decimal
from pathlib import Path

import pytest
from pydantic import BaseModel

from packages.llm.budget import DailyBudgetGuard, InMemoryCostCounter
from packages.llm.circuit import CircuitBreakerRegistry, CircuitState
from packages.llm.cost import InMemoryAICallRecorder, ModelRate, RateTable
from packages.llm.errors import (
    AllProvidersFailedError,
    DailyBudgetExceededError,
    ProviderBadRequestError,
    ProviderRateLimitedError,
    ProviderServerError,
    ProviderTimeoutError,
    SceneNotRoutedError,
    StructuredOutputError,
)
from packages.llm.gateway import LLMGateway, extract_json
from packages.llm.prompts import PromptRegistry
from packages.llm.providers.base import LLMProvider
from packages.llm.providers.zhipu import ZhipuGLMProvider
from packages.llm.routing import SceneRouter, build_scene_configs
from packages.llm.types import Message, ProviderRequest, ProviderResponse, Scene, Usage

RATE = ModelRate(
    input_cents_per_mtok=Decimal("100"),
    cached_input_cents_per_mtok=Decimal("10"),
    output_cents_per_mtok=Decimal("200"),
)


class FakeProvider:
    """Returns queued responses/exceptions in order; repeats the last forever."""

    def __init__(
        self,
        name: str,
        outcomes: Sequence[ProviderResponse | Exception],
        default_model: str = "fake-model",
    ) -> None:
        self._name = name
        self._outcomes = list(outcomes)
        self._default_model = default_model
        self.requests: list[ProviderRequest] = []
        self.closed = False

    @property
    def name(self) -> str:
        return self._name

    @property
    def default_model(self) -> str:
        return self._default_model

    async def complete(self, request: ProviderRequest) -> ProviderResponse:
        self.requests.append(request)
        outcome = self._outcomes[min(len(self.requests) - 1, len(self._outcomes) - 1)]
        if isinstance(outcome, Exception):
            raise outcome
        return outcome

    async def aclose(self) -> None:
        self.closed = True


def response(text: str, model: str = "fake-model") -> ProviderResponse:
    return ProviderResponse(
        text=text,
        model=model,
        usage=Usage(tokens_in=1000, tokens_out=200, tokens_cached=800),
        finish_reason="stop",
    )


def build_gateway(
    prompts_dir: Path,
    providers: Mapping[str, LLMProvider],
    *,
    primary: str = "p1",
    backup: str | None = "p2",
    budget_cents: str = "100000",
    budget_enabled: bool = True,
    failure_threshold: int = 3,
    structured_max_retries: int = 2,
    recorder: InMemoryAICallRecorder | None = None,
) -> tuple[LLMGateway, InMemoryAICallRecorder, DailyBudgetGuard]:
    recorder = recorder or InMemoryAICallRecorder()
    guard = DailyBudgetGuard(
        InMemoryCostCounter(),
        daily_budget_cents=Decimal(budget_cents),
        enabled=budget_enabled,
    )
    gateway = LLMGateway(
        providers=dict(providers),
        router=SceneRouter(
            build_scene_configs(
                primary_provider=primary,
                backup_provider=backup,
                structured_max_retries=structured_max_retries,
            )
        ),
        prompts=PromptRegistry(prompts_dir),
        rates=RateTable({"fake-model": RATE}, default=RATE),
        recorder=recorder,
        budget=guard,
        breakers=CircuitBreakerRegistry(failure_threshold=failure_threshold),
        id_factory=itertools.count(1).__next__,
    )
    return gateway, recorder, guard


# --------------------------------------------------------------------------
# routing
# --------------------------------------------------------------------------


async def test_routes_to_the_primary_provider(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("好的")])
    backup = FakeProvider("p2", [response("backup")])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})

    result = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "面试"})

    assert result.provider == "p1"
    assert result.text == "好的"
    assert backup.requests == []


async def test_scene_config_drives_model_and_sampling(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("ok")], default_model="deepseek-v4-flash")
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)

    await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})

    request = primary.requests[0]
    assert request.model == "deepseek-v4-flash"
    assert request.max_tokens == 800  # 架构文档 §7.2 输出预算
    assert request.timeout_seconds == 30.0


async def test_cacheable_system_prefix_is_the_first_message(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("ok")])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)

    await gateway.complete(
        scene=Scene.AGENT_CHAT,
        messages=[Message(role="user", content="历史消息")],
        prompt_vars={"user_input": "现在的问题"},
    )

    messages = primary.requests[0].messages
    assert messages[0].role == "system"
    assert "穿搭顾问助手" in messages[0].content
    assert messages[-1].content == "现在的问题"


async def test_unrouted_scene_raises(prompts_dir: Path) -> None:
    gateway = LLMGateway(
        providers={},
        router=SceneRouter({}),
        prompts=PromptRegistry(prompts_dir),
        rates=RateTable({}, default=RATE),
        recorder=InMemoryAICallRecorder(),
        budget=DailyBudgetGuard(
            InMemoryCostCounter(), daily_budget_cents=Decimal("1"), enabled=False
        ),
    )
    with pytest.raises(SceneNotRoutedError):
        await gateway.complete(scene=Scene.AGENT_CHAT)


async def test_unregistered_provider_is_reported_as_a_failure(prompts_dir: Path) -> None:
    gateway, _, _ = build_gateway(prompts_dir, {}, primary="ghost", backup=None)
    with pytest.raises(AllProvidersFailedError):
        await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})


# --------------------------------------------------------------------------
# failover
# --------------------------------------------------------------------------


@pytest.mark.parametrize(
    "failure",
    [
        ProviderTimeoutError("p1", "timed out"),
        ProviderServerError("p1", "502", status_code=502),
        ProviderRateLimitedError("p1", "429", status_code=429),
    ],
)
async def test_timeout_5xx_and_rate_limit_fail_over_to_the_backup(
    prompts_dir: Path, failure: Exception
) -> None:
    primary = FakeProvider("p1", [failure])
    backup = FakeProvider("p2", [response("备用回答")])
    gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})

    result = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})

    assert result.provider == "p2"
    assert result.text == "备用回答"
    # Both the failed and the successful attempt are attributable.
    assert [(r.provider, r.success) for r in recorder.records] == [("p1", False), ("p2", True)]


async def test_bad_request_does_not_fail_over(prompts_dir: Path) -> None:
    """A malformed request would fail on the backup too — burning money twice."""
    primary = FakeProvider("p1", [ProviderBadRequestError("p1", "400", status_code=400)])
    backup = FakeProvider("p2", [response("should not be used")])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})

    with pytest.raises(ProviderBadRequestError):
        await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
    assert backup.requests == []


async def test_all_providers_failing_raises_with_every_cause(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [ProviderTimeoutError("p1", "timeout")])
    backup = FakeProvider("p2", [ProviderServerError("p2", "500", status_code=500)])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})

    with pytest.raises(AllProvidersFailedError) as excinfo:
        await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
    assert len(excinfo.value.failures) == 2


async def test_zhipu_placeholder_participates_in_the_real_routing_path(
    prompts_dir: Path,
) -> None:
    """The backup slot is structural, but the code path through it is real."""
    primary = FakeProvider("p1", [ProviderTimeoutError("p1", "timeout")])
    gateway, recorder, _ = build_gateway(
        prompts_dir, {"p1": primary, "zhipu": ZhipuGLMProvider()}, backup="zhipu"
    )

    with pytest.raises(AllProvidersFailedError) as excinfo:
        await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})

    assert isinstance(excinfo.value.failures[1], NotImplementedError)
    assert {r.provider for r in recorder.records} == {"p1", "zhipu"}


# --------------------------------------------------------------------------
# circuit breaker
# --------------------------------------------------------------------------


async def test_open_circuit_skips_the_primary_entirely(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [ProviderServerError("p1", "500", status_code=500)])
    backup = FakeProvider("p2", [response("备用")])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup}, failure_threshold=1)

    first = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "a"})
    assert first.provider == "p2"
    assert gateway.breaker_states()["p1"] == CircuitState.OPEN.value

    second = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "b"})
    assert second.provider == "p2"
    assert len(primary.requests) == 1  # not attempted the second time


async def test_success_keeps_the_circuit_closed(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("ok")])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
    await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "a"})
    assert gateway.breaker_states()["p1"] == CircuitState.CLOSED.value


# --------------------------------------------------------------------------
# structured output
# --------------------------------------------------------------------------


class Outfit(BaseModel):
    top: str
    bottom: str


async def test_structured_output_is_parsed_and_validated(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response('{"top": "白衬衫", "bottom": "西裤"}')])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)

    result = await gateway.complete(
        scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
    )

    assert result.parsed_as(Outfit).top == "白衬衫"
    assert result.attempts == 1
    assert primary.requests[0].json_mode is True
    assert "JSON Schema" in primary.requests[0].messages[-1].content


async def test_invalid_structured_output_is_retried_with_the_validation_error(
    prompts_dir: Path,
) -> None:
    primary = FakeProvider(
        "p1",
        [
            response("这不是 JSON"),
            response('{"top": "针织衫"}'),  # missing `bottom`
            response('{"top": "针织衫", "bottom": "阔腿裤"}'),
        ],
    )
    gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)

    result = await gateway.complete(
        scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
    )

    assert result.attempts == 3
    assert result.parsed_as(Outfit).bottom == "阔腿裤"
    # Every attempt costs money and every attempt is recorded.
    assert len(recorder.records) == 3
    assert [r.success for r in recorder.records] == [False, False, True]
    assert recorder.records[0].error_code == "llm_structured_output_invalid"
    # The repair turn carries the validation error back to the model.
    assert "校验错误" in primary.requests[1].messages[-1].content


async def test_structured_output_gives_up_and_fails_over(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("never json")])
    backup = FakeProvider("p2", [response('{"top": "T恤", "bottom": "牛仔裤"}')])
    gateway, _, _ = build_gateway(
        prompts_dir, {"p1": primary, "p2": backup}, structured_max_retries=1
    )

    result = await gateway.complete(
        scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
    )

    assert result.provider == "p2"
    assert len(primary.requests) == 2  # 1 + structured_max_retries


async def test_structured_failure_on_the_only_provider_surfaces(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("never json")])
    gateway, _, _ = build_gateway(
        prompts_dir, {"p1": primary}, backup=None, structured_max_retries=0
    )
    with pytest.raises(AllProvidersFailedError) as excinfo:
        await gateway.complete(
            scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, schema=Outfit
        )
    assert isinstance(excinfo.value.failures[0], StructuredOutputError)


@pytest.mark.parametrize(
    ("raw", "expected"),
    [
        ('{"a": 1}', '{"a": 1}'),
        ('```json\n{"a": 1}\n```', '{"a": 1}'),
        ('好的：\n{"a": 1}\n以上。', '{"a": 1}'),
        ("no json here", "no json here"),
    ],
)
def test_extract_json_handles_fenced_and_chatty_replies(raw: str, expected: str) -> None:
    assert extract_json(raw) == expected


# --------------------------------------------------------------------------
# cost attribution + daily budget
# --------------------------------------------------------------------------


async def test_every_call_records_a_cost_row(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("ok")])
    gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)

    result = await gateway.complete(
        scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}, user_id=42, ref_id="req-1"
    )

    record = recorder.records[0]
    assert record.scene == "agent_chat"
    assert record.provider == "p1"
    assert record.model == "fake-model"
    assert record.prompt_id == "agent_chat"
    assert record.prompt_version == "v1"  # 可复现: 记录 prompt 版本
    assert record.user_id == 42
    assert record.ref_id == "req-1"
    assert record.usage.tokens_cached == 800
    assert record.cost_cents == result.cost_cents > Decimal("0")
    assert record.created_at is not None
    assert record.rate_tier  # 复核成本要知道当次单价档位


async def test_failed_calls_are_recorded_too(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [ProviderTimeoutError("p1", "timed out")])
    gateway, recorder, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)

    with pytest.raises(AllProvidersFailedError):
        await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})

    record = recorder.records[0]
    assert record.success is False
    assert record.error_code == "llm_timeout"
    assert record.cost_cents == Decimal("0")
    # 失败行也要带档位：排查「为什么这天成本高」时失败重试也是成本的一部分
    assert record.rate_tier


async def test_a_broken_recorder_never_breaks_the_user_request(prompts_dir: Path) -> None:
    class ExplodingRecorder:
        async def record(self, record: object) -> None:
            msg = "database is down"
            raise RuntimeError(msg)

    primary = FakeProvider("p1", [response("ok")])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary}, backup=None)
    gateway._recorder = ExplodingRecorder()

    result = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"})
    assert result.text == "ok"


async def test_spend_accumulates_into_the_daily_budget(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("ok")])
    gateway, _, guard = build_gateway(prompts_dir, {"p1": primary}, backup=None)

    first = await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "a"})
    await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "b"})

    assert await guard.spent_today() == first.cost_cents * 2


async def test_daily_budget_rejects_non_critical_scenes_before_spending(
    prompts_dir: Path,
) -> None:
    # block_all 档默认放行 agent_chat（纯文本对话），用非 allow 场景断言拒绝
    primary = FakeProvider("p1", [response("ok")])
    gateway, _, guard = build_gateway(
        prompts_dir, {"p1": primary}, backup=None, budget_cents="0.00001"
    )
    await guard.add_cost(Decimal("1"))

    with pytest.raises(DailyBudgetExceededError):
        await gateway.complete(
            scene=Scene.REPORT_COPY, prompt_vars={"payload": "诊断结论与规则命中"}
        )
    assert primary.requests == []  # rejected, not merely alerted


async def test_daily_budget_lets_critical_scenes_through(prompts_dir: Path) -> None:
    # 意图分类把守越界拒答，静默降级会让超范围提问漏过去，因此不受预算闸门限制
    primary = FakeProvider("p1", [response("scene_consult")])
    gateway, _, guard = build_gateway(
        prompts_dir, {"p1": primary}, backup=None, budget_cents="0.00001"
    )
    await guard.add_cost(Decimal("1"))

    result = await gateway.complete(
        scene=Scene.INTENT_CLASSIFY, prompt_vars={"user_input": "面试穿什么"}
    )
    assert result.text == "scene_consult"


async def test_aclose_closes_every_provider(prompts_dir: Path) -> None:
    primary = FakeProvider("p1", [response("ok")])
    backup = FakeProvider("p2", [response("ok")])
    gateway, _, _ = build_gateway(prompts_dir, {"p1": primary, "p2": backup})
    await gateway.aclose()
    assert primary.closed and backup.closed