deerflow-code/offline-backend-20260512/backend/tests/test_llm_error_handling_middleware.py

"""LLM 错误分类：httpx 超时族必须可重试（瞬时网络故障不能直接终结 run）。

线上实测：兼容模式流式响应中途 200 OK 后停发数据 ~17s，抛 ``httpx.ReadTimeout``。
旧名单只认 ``ReadError`` / ``RemoteProtocolError``，ReadTimeout 被归为 generic
（不可重试），1 次尝试即兜底报错、整个 run 终止。
"""

from __future__ import annotations

import asyncio
from types import SimpleNamespace

import httpx
import pytest
from langchain_core.messages import AIMessage, AIMessageChunk

from deerflow.agents.middlewares.llm_error_handling_middleware import LLM_ERROR_MARKER, LLMErrorHandlingMiddleware
from deerflow.models.patched_openai import OpenAICompatibleUpstreamError, PatchedChatOpenAI


def _middleware() -> LLMErrorHandlingMiddleware:
    config = SimpleNamespace(
        circuit_breaker=SimpleNamespace(failure_threshold=3, recovery_timeout_sec=30),
    )
    middleware = LLMErrorHandlingMiddleware(app_config=config)
    middleware.retry_base_delay_ms = 0
    middleware.retry_cap_delay_ms = 0
    return middleware


def _openai_model() -> PatchedChatOpenAI:
    return PatchedChatOpenAI(model="test-model", api_key="test-key")


def test_httpx_read_timeout_is_transient() -> None:
    assert _middleware()._classify_error(httpx.ReadTimeout("read timed out")) == (True, "transient")


def test_httpx_timeout_family_is_transient() -> None:
    middleware = _middleware()
    for exc in (
        httpx.ConnectTimeout("connect timed out"),
        httpx.WriteTimeout("write timed out"),
        httpx.PoolTimeout("pool timed out"),
        httpx.ConnectError("connection refused"),
        TimeoutError(),
    ):
        assert middleware._classify_error(exc) == (True, "transient"), exc


def test_value_error_stays_generic() -> None:
    assert _middleware()._classify_error(ValueError("boom")) == (False, "generic")


def test_openai_compatible_embedded_error_is_retriable_upstream() -> None:
    error = RuntimeError("OpenAI-compatible upstream application error (code=2056)")
    assert _middleware()._classify_error(error) == (True, "upstream")


def test_openai_compatible_null_choices_error_is_retriable_upstream() -> None:
    error = TypeError("Received response with null value for 'choices'.")
    assert _middleware()._classify_error(error) == (True, "upstream")


def test_insufficient_balance_stays_non_retriable_quota() -> None:
    """The 200-body compatibility fix must not hide a real billing failure."""

    class PaymentError(RuntimeError):
        status_code = 402

    error = PaymentError("Insufficient Balance")
    assert _middleware()._classify_error(error) == (False, "quota")


def test_minimax_token_plan_limit_stays_non_retriable_quota() -> None:
    error = OpenAICompatibleUpstreamError(
        code=2056,
        message="Token Plan limit reached",
    )

    assert _middleware()._classify_error(error) == (False, "quota")


def test_non_stream_embedded_application_error_is_raised() -> None:
    response = {
        "id": "failed-response",
        "object": "chat.completion",
        "created": 0,
        "model": "test-model",
        "choices": None,
        "base_resp": {
            "status_code": 2056,
            "status_msg": "Token Plan limit reached",
        },
    }

    with pytest.raises(OpenAICompatibleUpstreamError, match=r"code=2056") as exc_info:
        _openai_model()._create_chat_result(response)

    assert exc_info.value.code == 2056
    assert "Token Plan limit reached" in str(exc_info.value)


def test_stream_embedded_application_error_is_raised() -> None:
    chunk = {
        "id": "failed-response",
        "object": "chat.completion.chunk",
        "created": 0,
        "model": "test-model",
        "choices": [],
        "base_resp": {
            "status_code": 2056,
            "status_msg": "Token Plan limit reached",
        },
    }

    with pytest.raises(OpenAICompatibleUpstreamError, match=r"code=2056"):
        _openai_model()._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, None)


def test_normal_empty_usage_chunk_is_not_treated_as_error() -> None:
    chunk = {
        "id": "usage-response",
        "object": "chat.completion.chunk",
        "created": 0,
        "model": "test-model",
        "choices": [],
        "usage": {
            "prompt_tokens": 3,
            "completion_tokens": 2,
            "total_tokens": 5,
        },
    }

    result = _openai_model()._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, None)

    assert result is not None
    assert result.message.usage_metadata == {
        "input_tokens": 3,
        "output_tokens": 2,
        "total_tokens": 5,
        "input_token_details": {},
        "output_token_details": {},
    }


def test_read_timeout_is_retried_and_recovers() -> None:
    """第一次 ReadTimeout、第二次成功：必须重试并返回真实结果。"""
    middleware = _middleware()
    calls = {"count": 0}

    async def handler(_request):
        calls["count"] += 1
        if calls["count"] == 1:
            raise httpx.ReadTimeout("read timed out")
        return AIMessage(content="ok")

    result = asyncio.run(middleware.awrap_model_call(SimpleNamespace(), handler))

    assert calls["count"] == 2
    assert result.content == "ok"


def test_persistent_read_timeout_falls_back_with_transient_reason() -> None:
    """连续超时到重试上限：兜底消息 reason=transient（而非 generic）。"""
    middleware = _middleware()

    async def handler(_request):
        raise httpx.ReadTimeout("read timed out")

    result = asyncio.run(middleware.awrap_model_call(SimpleNamespace(), handler))

    assert isinstance(result, AIMessage)
    assert result.additional_kwargs[LLM_ERROR_MARKER] == "transient"