vent_agent_2.0/tests/test_model_retries.py

"""使用真实模型客户端与离线 HTTP 响应验证限流恢复和重试边界。"""
import asyncio
from contextlib import asynccontextmanager
import json

import httpx
import openai._base_client as openai_base
import pytest
from langchain_openai.chat_models.base import (
    OpenAIAuthenticationError, OpenAIConnectionError, OpenAIRateLimitError,
)

from agents import model_factory

MODEL = "Qwen/Qwen3.5-35B-A3B-FP8"


def limited():
    return httpx.Response(429, headers={"Retry-After": "11"}, json={
        "error": {"code": "concurrent_limit", "message": "并发数超限，请稍后重试",
                  "type": "rate_limit_error"}})


def sse():
    def chunk(delta, finish=None):
        return "data: " + json.dumps({
            "id": "chatcmpl-retry-test", "object": "chat.completion.chunk",
            "created": 0, "model": MODEL,
            "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
        }, ensure_ascii=False) + "\n\n"
    return (chunk({"role": "assistant", "content": "已恢复"})
            + chunk({}, "stop") + "data: [DONE]\n\n").encode()


@asynccontextmanager
async def mocked_model(monkeypatch, handler, retries=8):
    monkeypatch.setattr(model_factory.config, "OPENAI_MAX_RETRIES", retries)
    instantiate = model_factory._instantiate
    with httpx.Client(transport=httpx.MockTransport(handler)) as sync:
        async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
            monkeypatch.setattr(model_factory, "_instantiate", lambda name, options: instantiate(
                name, {**options, "http_client": sync, "http_async_client": client}))
            yield model_factory.build_model(
                MODEL, base_url="https://model.invalid/v1", api_key="sk-test")


@pytest.mark.asyncio
async def test_stream_recovers_after_three_concurrent_limits_and_honors_retry_after(monkeypatch):
    requests, waits = [], []

    def handler(request):
        requests.append(request.content)
        if len(requests) <= 3:
            return limited()
        return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=sse())

    async def sleep(delay):
        waits.append(delay)

    monkeypatch.setattr(openai_base.anyio, "sleep", sleep)
    async with mocked_model(monkeypatch, handler) as model:
        chunks = [chunk async for chunk in model.astream("测试")]
    assert "".join(chunk.content for chunk in chunks) == "已恢复"
    assert len(requests) == 4 and len(set(requests)) == 1
    assert waits == [11, 11, 11]


@pytest.mark.asyncio
@pytest.mark.parametrize("retries", [0, 2])
async def test_retry_budget_is_bounded_and_zero_disables_it(monkeypatch, retries):
    calls, waits = [], []

    def handler(request):
        calls.append(request)
        return limited()

    async def sleep(delay):
        waits.append(delay)

    monkeypatch.setattr(openai_base.anyio, "sleep", sleep)
    async with mocked_model(monkeypatch, handler, retries) as model:
        with pytest.raises(OpenAIRateLimitError, match="concurrent_limit"):
            await model.ainvoke("测试")
    assert len(calls) == retries + 1
    assert waits == [11] * retries


@pytest.mark.asyncio
async def test_authentication_failure_is_not_retried(monkeypatch):
    calls = []

    def handler(request):
        calls.append(request)
        return httpx.Response(401, json={"error": {"message": "invalid key", "type": "authentication_error"}})

    async with mocked_model(monkeypatch, handler) as model:
        with pytest.raises(OpenAIAuthenticationError):
            await model.ainvoke("测试")
    assert len(calls) == 1


@pytest.mark.asyncio
async def test_cancelling_during_retry_wait_stops_further_requests(monkeypatch):
    calls = []
    waiting = asyncio.Event()

    def handler(request):
        calls.append(request)
        return limited()

    async def sleep(delay):
        waiting.set()
        await asyncio.Event().wait()

    monkeypatch.setattr(openai_base.anyio, "sleep", sleep)
    async with mocked_model(monkeypatch, handler) as model:
        task = asyncio.create_task(model.ainvoke("测试"))
        try:
            await asyncio.wait_for(waiting.wait(), timeout=5)
        finally:
            task.cancel()
            with pytest.raises(asyncio.CancelledError):
                await task
    assert len(calls) == 1


@pytest.mark.asyncio
async def test_broken_stream_is_not_replayed(monkeypatch):
    calls, received = [], []

    class BrokenStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            yield sse().split(b"data: [DONE]")[0]
            raise httpx.ReadError("stream interrupted")

    def handler(request):
        calls.append(request)
        return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, stream=BrokenStream())

    async with mocked_model(monkeypatch, handler) as model:
        with pytest.raises(OpenAIConnectionError):
            async for chunk in model.astream("测试"):
                received.append(chunk.content)
    assert "".join(received) == "已恢复"
    assert len(calls) == 1


def test_negative_retry_budget_is_rejected(monkeypatch):
    monkeypatch.setattr(model_factory.config, "OPENAI_MAX_RETRIES", -1)
    with pytest.raises(ValueError, match="OPENAI_MAX_RETRIES"):
        model_factory.build_model(MODEL)


@pytest.mark.asyncio
async def test_http_200_empty_sse_does_not_become_a_successful_execution(monkeypatch, tmp_path, pricing_ready):
    from langchain.agents import create_agent
    from langgraph.checkpoint.memory import InMemorySaver
    from test_execution import prepared, assert_released

    monkeypatch.setattr(model_factory.config, 'ROOT', tmp_path)
    empty = sse().replace('已恢复'.encode(), b'')
    calls = []

    def handler(request):
        calls.append(request)
        return httpx.Response(200, headers={'Content-Type': 'text/event-stream'}, content=empty)

    async with mocked_model(monkeypatch, handler) as model:
        graph = create_agent(model=model, checkpointer=InMemorySaver())
        task = await prepared(graph)
        result = await task.run()

    assert result.status == 'failed' and result.error == 'empty_model_response'
    assert result.output == '' and result.structured_output is None
    assert result.usage['usage_source'] == 'estimate'
    assert result.diagnostics['model_responses'] == 1
    assert result.diagnostics['last_response']['finish_reason'] == 'stop'
    assert len(calls) == 1, '已建立并正常结束的 HTTP 200 流不会触发 SDK 的 HTTP 重试'
    assert_released()