PolyGateway/tests/unit/test_retry.py

"""RetryMW 尝试编排测试(保真蓝本 CHS governance.py:107-268)。

用真实内存后端 + FakeClock + 可编程 fake transport,验证换源/退避/熔断
写回/permit 结算/取消穿透等治理行为。
"""

import asyncio

import pytest

from polygateway.backends.memory.breaker import InMemoryGate
from polygateway.backends.memory.limiter import InMemoryLimiter
from polygateway.errors import (
    AllSourcesExhausted,
    CircuitOpenError,
    RequestRejectedError,
    ResultInvalidError,
    SourceDeadError,
    TransientError,
)
from polygateway.middleware.retry import RetryMW
from polygateway.sources import RoundRobinSelector, SourceCooldownMemo
from polygateway.types import (
    BackpressurePolicy,
    BreakerConfig,
    ChatRequest,
    GlobalLimits,
    RetryPolicy,
    SourceConfig,
    TransportResult,
)
from tests.contracts.conftest import FakeClock

_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)


def _src(name, **overrides):
    base = {
        "name": name,
        "provider": "openai",
        "base_url": "https://gw.example/v1",
        "api_key": "sk",
        "model": "m",
        "timeout_s": 10.0,
    }
    base.update(overrides)
    return SourceConfig(**base)


def _ok(content="ok"):
    return TransportResult(
        content=content,
        thinking="",
        prompt_tokens=10,
        completion_tokens=5,
        usage_source="measured",
        ttft_ms=12.0,
        max_inter_token_ms=3.0,
        raw={},
    )


class FakeTransport:
    """按脚本逐次返回结果或抛异常;记录每次 (source_name, call_id)。"""

    def __init__(self, script):
        self.script = list(script)
        self.calls = []

    async def complete(self, *, messages, source, stream, overlay, call_id):
        self.calls.append((source.name, call_id))
        action = self.script.pop(0)
        if isinstance(action, Exception):
            raise action
        if action == "hang":
            await asyncio.Event().wait()
        return action


class FakeSleep:
    """记录退避时长,立即返回(不真等)。"""

    def __init__(self):
        self.delays = []

    async def __call__(self, seconds):
        self.delays.append(seconds)


def _harness(
    sources,
    script,
    *,
    clock=None,
    max_attempts=3,
    quota_full="wait",
    global_limits=_NO_GLOBAL,
    rng=lambda: 0.0,
    selector=None,
):
    clock = clock or FakeClock()
    limiter = InMemoryLimiter(
        scope="llm",
        sources={s.name: s for s in sources},
        global_limits=global_limits,
        lease_ttl_s=100.0,
        now=clock,
    )
    gate = InMemoryGate(config=_BREAKER, now=clock)
    transport = FakeTransport(script)
    sleep = FakeSleep()
    mw = RetryMW(
        scope="llm",
        sources=sources,
        selector=selector if selector is not None else RoundRobinSelector(),
        limiter=limiter,
        gate=gate,
        transport=transport,
        retry=RetryPolicy(max_attempts=max_attempts, backoff_base_s=2.0, backoff_max_s=30.0),
        backpressure=BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.01),
        quota_full=quota_full,
        cooldown_memo=SourceCooldownMemo(now=clock),
        emitter=None,
        now=clock,
        sleep=sleep,
        rng=rng,
    )
    return mw, limiter, gate, transport, sleep, clock


_REQ = ChatRequest(messages=[{"role": "user", "content": "hi"}])


class TestSuccessPath:
    async def test_first_attempt_success_builds_response(self):
        mw, limiter, gate, transport, sleep, _ = _harness([_src("a")], [_ok("hello")])
        resp = await mw(_REQ)
        assert resp.content == "hello"
        assert resp.source_name == "a" and resp.provider == "openai"
        assert resp.cache_hit is False and resp.call_id
        assert (await limiter.source_stats("a")).inflight == 0  # permit 已释放
        assert await limiter.progress_age_s() < 5.0  # mark_progress 已调用
        assert sleep.delays == []

    async def test_settle_uses_actual_usage(self):
        src = _src("a", tpm=1000, est_tokens=400)
        mw, limiter, *_ = _harness([src], [_ok()])
        await mw(_REQ)
        # 预扣 400,实际 15 → settle 后窗口只记 15
        assert (await limiter.source_stats("a")).tpm_used == 15


class TestRetryAndFailover:
    async def test_transient_switches_source_then_succeeds(self):
        mw, _, _, transport, sleep, _ = _harness(
            [_src("a"), _src("b")], [TransientError("boom"), _ok()]
        )
        resp = await mw(_REQ)
        assert [name for name, _ in transport.calls] == ["a", "b"]
        assert resp.source_name == "b"
        assert len(sleep.delays) == 1  # 瞬时错误退避一次

    async def test_each_attempt_gets_fresh_call_id(self):
        mw, _, _, transport, _, _ = _harness([_src("a")], [TransientError("x"), _ok()])
        await mw(_REQ)
        ids = [cid for _, cid in transport.calls]
        assert len(ids) == 2 and ids[0] != ids[1]

    async def test_max_attempts_is_total_attempts(self):
        mw, _, _, transport, _, _ = _harness(
            [_src("a")],
            [TransientError("1"), TransientError("2"), TransientError("3")],
            max_attempts=3,
        )
        with pytest.raises(AllSourcesExhausted) as ei:
            await mw(_REQ)
        assert len(transport.calls) == 3  # 恰 3 次总尝试(含首次)
        assert ei.value.reason == "retry_exhausted"
        assert ei.value.per_source_reasons.get("a") == "network_error"

    async def test_backoff_formula_and_retry_after_max(self):
        # rng=0 → jitter 因子 0.5;第一次退避 = 2*2^0*0.5 = 1.0
        mw, _, _, _, sleep, _ = _harness([_src("a")], [TransientError("x"), _ok()])
        await mw(_REQ)
        assert sleep.delays == [1.0]
        # Retry-After 提示更大时取提示值
        mw2, _, _, _, sleep2, _ = _harness(
            [_src("a")], [TransientError("x", retry_after_s=7.5), _ok()]
        )
        await mw2(_REQ)
        assert sleep2.delays == [7.5]

    async def test_source_dead_switches_immediately_and_force_opens(self):
        mw, _, gate, transport, sleep, _ = _harness(
            [_src("a"), _src("b")], [SourceDeadError("401"), _ok()]
        )
        resp = await mw(_REQ)
        assert resp.source_name == "b"
        assert sleep.delays == []  # 源死亡不退避
        assert not (await gate.try_enter("a", "w")).allowed  # a 已 force_open


class TestNonRetryableOutcomes:
    async def test_request_rejected_propagates_without_retry(self):
        exc = RequestRejectedError("400", source_name="a", status_code=400)
        mw, _, gate, transport, _, _ = _harness([_src("a")], [exc])
        with pytest.raises(RequestRejectedError):
            await mw(_REQ)
        assert len(transport.calls) == 1
        # 网关已应答 → 记成功,熔断计数未增长
        assert (await gate.try_enter("a", "w")).allowed

    async def test_result_invalid_records_success_and_propagates(self):
        mw, limiter, gate, transport, _, _ = _harness(
            [_src("a")], [ResultInvalidError("bad json", raw_text="{oops")]
        )
        with pytest.raises(ResultInvalidError):
            await mw(_REQ)
        assert len(transport.calls) == 1  # 坏结果不重试
        assert (await gate.try_enter("a", "w")).allowed  # 熔断记成功
        assert (await limiter.source_stats("a")).inflight == 0


class TestScopeUnavailable:
    async def test_all_sources_circuit_open(self):
        clock = FakeClock()
        script = [TransientError(str(i)) for i in range(9)]
        mw, _, gate, _, _, _ = _harness([_src("a")], script, clock=clock, max_attempts=99)
        # 3 次失败后 a 开路 → 第 4 次尝试选不到源且 gate_rejections==全部 → CircuitOpen
        with pytest.raises(CircuitOpenError) as ei:
            await mw(_REQ)
        assert ei.value.reason == "circuit_open"
        assert ei.value.retry_after_s > 0
        assert ei.value.per_source_reasons  # 携逐源原因

    async def test_no_sources_configured(self):
        mw, *_ = _harness([], [])
        with pytest.raises(AllSourcesExhausted) as ei:
            await mw(_REQ)
        assert ei.value.reason == "no_sources"

    async def test_quota_fail_fast(self):
        src = _src("a", max_concurrency=1)
        mw, limiter, _, _, _, _ = _harness([src], [_ok()], quota_full="fail_fast")
        held = await limiter.try_acquire("a", 0)  # 外部占满并发
        assert held is not None
        with pytest.raises(AllSourcesExhausted) as ei:
            await mw(_REQ)
        assert ei.value.reason == "quota_exhausted"

    async def test_quota_wait_polls_until_slot_frees(self):
        src = _src("a", max_concurrency=1)
        clock = FakeClock()
        limiter = InMemoryLimiter(
            scope="llm",
            sources={"a": src},
            global_limits=_NO_GLOBAL,
            lease_ttl_s=100.0,
            now=clock,
        )
        held = await limiter.try_acquire("a", 0)
        released = {"done": False}

        async def sleep_and_release(seconds):
            if not released["done"]:
                released["done"] = True
                await held.release()

        gate = InMemoryGate(config=_BREAKER, now=clock)
        transport = FakeTransport([_ok()])
        mw = RetryMW(
            scope="llm",
            sources=[src],
            selector=RoundRobinSelector(),
            limiter=limiter,
            gate=gate,
            transport=transport,
            retry=RetryPolicy(max_attempts=3, backoff_base_s=2.0, backoff_max_s=30.0),
            backpressure=BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.01),
            quota_full="wait",
            cooldown_memo=SourceCooldownMemo(now=clock),
            emitter=None,
            now=clock,
            sleep=sleep_and_release,
            rng=lambda: 0.0,
        )
        resp = await mw(_REQ)
        assert resp.content == "ok" and released["done"]


class TestCancellation:
    async def test_cancel_mid_flight_releases_permit(self):
        mw, limiter, _, _, _, _ = _harness([_src("a", max_concurrency=1)], ["hang"])
        task = asyncio.ensure_future(mw(_REQ))
        await asyncio.sleep(0.05)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert (await limiter.source_stats("a")).inflight == 0  # finally 释放

    async def test_cancel_mid_backoff_propagates_with_no_held_permit(self):
        """退避 sleep 中取消: CancelledError 穿透,且 permit 早已在 finally 释放。"""
        clock = FakeClock()
        src = _src("a", max_concurrency=1)
        limiter = InMemoryLimiter(
            scope="llm",
            sources={"a": src},
            global_limits=_NO_GLOBAL,
            lease_ttl_s=100.0,
            now=clock,
        )
        mw = RetryMW(
            scope="llm",
            sources=[src],
            selector=RoundRobinSelector(),
            limiter=limiter,
            gate=InMemoryGate(config=_BREAKER, now=clock),
            transport=FakeTransport([TransientError("x"), _ok()]),
            retry=RetryPolicy(max_attempts=3, backoff_base_s=30.0, backoff_max_s=60.0),
            backpressure=BackpressurePolicy(300.0, 0.01),
            cooldown_memo=SourceCooldownMemo(now=clock),
            emitter=None,
            now=clock,
            sleep=asyncio.sleep,
            rng=lambda: 0.5,
        )
        task = asyncio.ensure_future(mw(_REQ))
        await asyncio.sleep(0.05)  # 第一次失败后进入 30s 真实退避
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert (await limiter.source_stats("a")).inflight == 0  # 退避期不占并发槽

    async def test_cancel_probe_releases_probe_lease(self):
        clock = FakeClock()
        mw, _, gate, _, _, _ = _harness(
            [_src("a")],
            [TransientError("1"), TransientError("2"), TransientError("3"), "hang"],
            clock=clock,
            max_attempts=99,
        )
        # 三连失败开路
        with pytest.raises(CircuitOpenError):
            await mw(_REQ)
        clock.advance(_BREAKER.cooldown_s + 1)
        task = asyncio.ensure_future(mw(_REQ))  # 半开探针 → hang
        await asyncio.sleep(0.05)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        # 探针租约已归还: 下一 caller 立即拿到探针而非等租约过期
        nxt = await gate.try_enter("a", "w2")
        assert nxt.allowed and nxt.is_probe


class StaticSelector:
    """固定配置序,隔离测试调用内降权(不带 record_outcome)。"""

    def order(self, sources, stats):
        return list(sources)


class RecordingSelector(StaticSelector):
    def __init__(self):
        self.outcomes = []

    def record_outcome(self, source_name, ok):
        self.outcomes.append((source_name, ok))


class ExplodingSelector(StaticSelector):
    def record_outcome(self, source_name, ok):
        raise RuntimeError("sink boom")


class TestM25Orchestration:
    """M2.5 设计 §3.3: 调用内失败降权 + 健康喂数(先红后绿)。"""

    async def test_failed_source_demoted_after_two_strikes(self):
        # 失败 1 次仍首选(原地退避重试);失败 2 次让位次优源
        mw, _, _, transport, _, _ = _harness(
            [_src("a"), _src("b")],
            [TransientError("1"), TransientError("2"), _ok()],
            selector=StaticSelector(),
        )
        resp = await mw(_REQ)
        assert [n for n, _ in transport.calls] == ["a", "a", "b"]
        assert resp.source_name == "b"

    async def test_attempt_fails_reset_between_calls(self):
        mw, _, _, transport, _, _ = _harness(
            [_src("a"), _src("b")],
            [TransientError("1"), TransientError("2"), _ok(), _ok()],
            selector=StaticSelector(),
        )
        await mw(_REQ)
        await mw(_REQ)  # 新调用状态清零: 回到首选 a
        assert [n for n, _ in transport.calls] == ["a", "a", "b", "a"]

    async def test_outcome_feeding_success_and_transient(self):
        sel = RecordingSelector()
        mw, _, _, _, _, _ = _harness(
            [_src("a"), _src("b")], [TransientError("1"), _ok()], selector=sel
        )
        await mw(_REQ)
        assert sel.outcomes == [("a", False), ("a", True)]

    async def test_outcome_feeding_source_dead(self):
        sel = RecordingSelector()
        mw, _, _, _, _, _ = _harness(
            [_src("a"), _src("b")], [SourceDeadError("401"), _ok()], selector=sel
        )
        await mw(_REQ)
        assert sel.outcomes == [("a", False), ("b", True)]

    async def test_result_invalid_and_rejected_not_fed(self):
        sel = RecordingSelector()
        mw, _, _, _, _, _ = _harness(
            [_src("a")], [ResultInvalidError("bad", raw_text="x")], selector=sel
        )
        with pytest.raises(ResultInvalidError):
            await mw(_REQ)
        sel2 = RecordingSelector()
        mw2, _, _, _, _, _ = _harness(
            [_src("a")],
            [RequestRejectedError("400", source_name="a", status_code=400)],
            selector=sel2,
        )
        with pytest.raises(RequestRejectedError):
            await mw2(_REQ)
        assert sel.outcomes == [] and sel2.outcomes == []

    async def test_outcome_sink_exception_swallowed(self):
        mw, _, _, _, _, _ = _harness([_src("a")], [_ok()], selector=ExplodingSelector())
        resp = await mw(_REQ)
        assert resp.content == "ok"  # 喂数异常不得打断真实成功返回

    async def test_result_invalid_gate_window_untouched(self):
        # 坏结果 ≠ 坏服务: count_attempt=False,失败率窗口 attempts 不得增长
        mw, _, gate, _, _, _ = _harness([_src("a")], [ResultInvalidError("bad", raw_text="x")])
        with pytest.raises(ResultInvalidError):
            await mw(_REQ)
        g = gate._gates["a"]
        assert g.a0 + g.a1 == 0