PolyGateway/tests/unit/test_stall_hedge.py

"""中段沉默对冲编排测试(1.3.8 issue #28 设计 §5.2 / 计划批次 D)。

设施纪律(计划 §5): 复用 test_hedge.py 范式——两源 scope、事件驱动假 transport、
**真实 loop 钟**(不注入 FakeClock 驱动触发;注入钟仅用于赢家净口径精确断言,
触发路径结构性不读它)、阈值 0.05s、断言容差 4–10×;取消窗口用 entered Event
栅栏钉死,禁 sleep 撞窗口;不标 slow。

命名区隔(审查 N5): 本文件的 `stall_after_s` 是**对冲触发阈值**(相邻增量 chunk
间隔超过它 → 触发对冲);test_hedge.py:150 的 `stall_window_s` 是 **StallClock
预算**(非生产性等待的调用级上限)——同名词不同概念,勿混。
"""

import asyncio
import dataclasses
import time

import pytest

from polygateway.backends.memory.breaker import InMemoryGate
from polygateway.backends.memory.limiter import InMemoryLimiter
from polygateway.errors import AllSourcesExhausted, TransientError
from polygateway.middleware.retry import RetryMW
from polygateway.sources import SourceCooldownMemo
from polygateway.types import (
    BackpressurePolicy,
    BreakerConfig,
    ChatRequest,
    GlobalLimits,
    RetryPolicy,
    _CallContext,
)
from tests.contracts.conftest import FakeClock
from tests.unit.test_hedge import RecordingEmitter
from tests.unit.test_retry import StaticSelector, _ok, _src

_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
_STALL_AFTER_S = 0.05  # 真实 loop 钟对冲触发阈值(≠ StallClock 预算 stall_window_s,N5)


class StallHedgeTransport:
    """事件驱动假 transport(中段沉默剧本): 逐步骤控制 chunk 事件/挂起/成败。

    剧本 = 每源一串**步骤**,每次 `complete` 调用从上次位置继续执行到终态步骤
    (`ok`/`fail`)为止;多轮调用(重试)顺序消费后续步骤,耗尽即 AssertionError。

    步骤:
      ("chunk",)                      — 置 first_token_event 与 content_event(首个增量 chunk)
      ("chunk_after", delay)          — 睡 delay 后只置 content_event(后续增量 chunk)
      ("hang",)                       — 置位该源 entered,挂起直到该源 release(或被取消)
      ("ok", content)                 — 终态: 立即成功(带已登记的 stall/clock 效果)
      ("ok_after", delay, content)    — 终态: 真实 loop 钟睡 delay 后成功
      ("fail_after", delay, factory)  — 终态: 睡 delay 后抛 `factory()` 新造的异常
      ("stall", ms, count)            — 非终态: 给本剧本终态结果打自报 stall 值(假
                                        transport 自己报数,不经真实计时,计划 §3.3)
      ("advance", secs)               — 非终态: 推进注入钟(仅配 FakeClock 的用例)
    """

    def __init__(self, scripts: dict[str, list[tuple]], clock: FakeClock | None = None):
        self._scripts = {name: list(steps) for name, steps in scripts.items()}
        self._pos = dict.fromkeys(scripts, 0)
        self._clock = clock
        self.calls: list[str] = []
        # 逐次记录收到的 first_token_event/content_event 身份(对冲路恒 None 断言原料)
        self.ft_events: list[object] = []
        self.ce_events: list[object] = []
        self.entered = {name: asyncio.Event() for name in scripts}
        self.release = {name: asyncio.Event() for name in scripts}

    async def complete(
        self,
        *,
        messages,
        source,
        stream,
        overlay,
        call_id,
        reasoning_effort,
        first_token_event,
        content_event,
        response_mode,
    ):
        self.calls.append(source.name)
        self.ft_events.append(first_token_event)
        self.ce_events.append(content_event)
        steps = self._scripts[source.name]
        stall: tuple[float, float] | None = None
        while self._pos[source.name] < len(steps):
            action = steps[self._pos[source.name]]
            self._pos[source.name] += 1
            kind = action[0]
            if kind == "chunk":
                if first_token_event is not None:
                    first_token_event.set()
                if content_event is not None:
                    content_event.set()
            elif kind == "chunk_after":
                await asyncio.sleep(action[1])
                if content_event is not None:
                    content_event.set()
            elif kind == "hang":
                self.entered[source.name].set()
                await self.release[source.name].wait()
            elif kind == "stall":
                stall = (action[1], action[2])
            elif kind == "advance":
                assert self._clock is not None, "advance 步骤需要构造时注入 FakeClock"
                self._clock.advance(action[1])
            elif kind == "ok":
                result = _ok(action[1])
                if stall is not None:
                    result = dataclasses.replace(result, stall_ms=stall[0], stall_count=stall[1])
                return result
            elif kind == "ok_after":
                await asyncio.sleep(action[1])
                result = _ok(action[2])
                if stall is not None:
                    result = dataclasses.replace(result, stall_ms=stall[0], stall_count=stall[1])
                return result
            elif kind == "fail_after":
                await asyncio.sleep(action[1])
                raise action[2]()
            else:
                raise AssertionError(f"未知剧本步骤: {action!r}")
        raise AssertionError(f"{source.name} 剧本耗尽且无终态步骤")


def _harness(
    sources,
    transport,
    *,
    hedge_after_s=None,
    stall_after_s=None,
    max_attempts=3,
    emitter=None,
    selector=None,
    limiter=None,
    gate=None,
    stall_window_s=300.0,
    now=None,
):
    """真实 loop 钟装配(同 test_hedge.py 范式;`now` 仅供净口径精确断言注入 FakeClock)。

    `stall_window_s` 是 StallClock 预算(N5 命名区隔),与 `stall_after_s` 无关。
    """
    limiter = limiter or InMemoryLimiter(
        scope="llm",
        sources={s.name: s for s in sources},
        global_limits=_NO_GLOBAL,
        lease_ttl_s=100.0,
    )
    gate = gate or InMemoryGate(config=_BREAKER)
    mw = RetryMW(
        scope="llm",
        sources=sources,
        # 固定配置序: 原路恒为 s1、对冲路恒为 s2,断言不依赖选源器内部状态
        selector=selector if selector is not None else StaticSelector(),
        limiter=limiter,
        gate=gate,
        transport=transport,
        retry=RetryPolicy(max_attempts=max_attempts, backoff_base_s=0.01, backoff_max_s=0.05),
        backpressure=BackpressurePolicy(stall_window_s=stall_window_s, poll_interval_s=0.01),
        quota_full="wait",
        cooldown_memo=SourceCooldownMemo(),
        emitter=emitter,
        hedge_after_s=hedge_after_s,
        stall_after_s=stall_after_s,
        **({"now": now} if now is not None else {}),
    )
    return mw, limiter, gate


def _req(*, stream=True, ctx=None):
    """中段沉默是流式概念,本文件默认 stream=True(test_hedge.py 默认 False 相反)。"""
    return ChatRequest(
        messages=[{"role": "user", "content": "hi"}], stream=stream, call_context=ctx
    )


def _ctx(now=None):
    return _CallContext(now=now if now is not None else time.monotonic)


class TestStallHedgeTrigger:
    """阶段 2 触发(设计 §5.2): 首 token 已至后相邻增量 chunk 间隔超阈值 → 对冲。"""

    async def test_stall_triggers_hedge_after_first_token(self):
        """s1 首 chunk 后沉默、s2 即时成功: stall 窗触发对冲,赢家为对冲路。

        净口径顺延断言(计划批次 D 首行): 注入 FakeClock + s2 自报 stall
        (1500ms/2 段)、advance 2.0s → 赢家 `generation_ms == 500` 精确值,
        `generation_ms + stall_ms == total_latency_ms`(同钟还原公式)。
        遥测行(T3 缺口补齐钉): 赢家行 stall 两值为自报值,输家行恒 None。
        """
        clock = FakeClock()
        emitter = RecordingEmitter()
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {
                "s1": [("chunk",), ("hang",), ("ok", "late")],
                "s2": [("stall", 1500.0, 2), ("advance", 2.0), ("ok", "fast")],
            },
            clock=clock,
        )
        mw, _, _ = _harness(
            [s1, s2], transport, stall_after_s=_STALL_AFTER_S, emitter=emitter, now=clock
        )
        ctx = _ctx(now=clock)
        started = time.monotonic()
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        elapsed = time.monotonic() - started
        assert resp.content == "fast" and resp.source_name == "s2"
        assert elapsed < 10 * _STALL_AFTER_S  # 沉默路被对冲截断,而非等到看门狗
        assert transport.calls == ["s1", "s2"]
        stats = ctx.snapshot()
        assert stats.hedges == 1 and stats.hedge_won is True and stats.attempts == 2
        # 赢家路净口径: 墙钟 2000(advance 精确推进)− 沉默 1500 = 500
        assert stats.generation_ms == 500
        assert stats.stall_ms == 1500 and stats.stall_count == 2
        assert stats.generation_ms + stats.stall_ms == stats.total_latency_ms  # 还原公式
        # 逐次遥测: 赢家行带 stall 归因,被取消输家行恒 None(设计 §11-B2)
        winner = next(r for r in emitter.rows if r["source"] == "s2")
        loser = next(r for r in emitter.rows if r["source"] == "s1")
        assert winner["stall_ms"] == 1500 and winner["stall_count"] == 2
        assert winner["error"] is None
        assert loser["error"] == "hedge_cancelled"
        assert loser["stall_ms"] is None and loser["stall_count"] is None

    async def test_single_leg_attempt_row_carries_stall(self):
        """T3 缺口补齐钉: 未启用对冲的 chat 成功行也须经 `_emit` 落 stall 两值。

        计划 §3.4 指定 `retry._emit` 显式转发(归 T3);批次 C 落行断言只覆盖了
        emitter 直调,chat 路径漏传无人把守——本用例钉住端到端路径。
        """
        clock = FakeClock()
        emitter = RecordingEmitter()
        transport = StallHedgeTransport(
            {"s1": [("stall", 300.0, 1), ("advance", 1.0), ("ok", "x")]}, clock=clock
        )
        mw, _, _ = _harness([_src("s1")], transport, emitter=emitter, now=clock)
        ctx = _ctx(now=clock)
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert resp.content == "x"
        assert len(emitter.rows) == 1
        row = emitter.rows[0]
        assert row["error"] is None
        assert row["stall_ms"] == 300 and row["stall_count"] == 1
        stats = ctx.snapshot()
        assert stats.generation_ms == 700 and stats.stall_ms == 300

    async def test_stall_recheck_prevents_false_trigger(self):
        """chunk 在超时到期前一刻到达(0.049s 置位,阈值 0.05s)→ 不触发(N3① 复检)。

        只断"不触发 + attempts==1"结果,不断精确时序(计划 §7 抖动纪律)。
        """
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {
                "s1": [("chunk",), ("chunk_after", 0.049), ("ok", "done")],
                "s2": [("ok", "other")],
            }
        )
        mw, _, _ = _harness([s1, s2], transport, stall_after_s=_STALL_AFTER_S)
        ctx = _ctx()
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert resp.content == "done" and resp.source_name == "s1"
        assert transport.calls == ["s1"]  # 对冲从未发出
        stats = ctx.snapshot()
        assert stats.hedges == 0 and stats.attempts == 1 and stats.hedge_won is False

    async def test_stall_and_after_s_independent_matrix(self):
        """参数化四格(设计 §5.2 逐字): 两阈值独立生效,各格触发语义不串。"""
        # 格 1: 仅 AFTER_S + 流式 pre-TTFT 挂起 → AFTER_S 触发(1.3.7 语义不变)
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport({"s1": [("hang",), ("ok", "late")], "s2": [("ok", "fast")]})
        mw, _, _ = _harness([s1, s2], transport, hedge_after_s=_STALL_AFTER_S)
        ctx = _ctx()
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert resp.source_name == "s2"
        assert ctx.snapshot().hedges == 1

        # 格 2: 仅 STALL_AFTER_S + 流式 pre-TTFT 沉默 → **不触发**(阶段 1 不开窗)
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {"s1": [("hang",), ("ok", "s1ok")], "s2": [("ok", "other")]}
        )
        mw, _, _ = _harness([s1, s2], transport, stall_after_s=_STALL_AFTER_S)
        ctx = _ctx()
        task = asyncio.ensure_future(mw(_req(ctx=ctx)))
        await transport.entered["s1"].wait()
        await asyncio.sleep(4 * _STALL_AFTER_S)  # 4× 余量确认 stall 窗管不到 pre-TTFT
        assert transport.calls == ["s1"]
        transport.release["s1"].set()
        resp = await asyncio.wait_for(task, timeout=5)
        assert resp.source_name == "s1"
        stats = ctx.snapshot()
        assert stats.hedges == 0 and stats.attempts == 1

        # 格 3: 两者皆设 + 首 chunk 后沉默 → 阶段 2 stall 窗触发
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {"s1": [("chunk",), ("hang",), ("ok", "late")], "s2": [("ok", "fast")]}
        )
        mw, _, _ = _harness(
            [s1, s2], transport, hedge_after_s=_STALL_AFTER_S, stall_after_s=_STALL_AFTER_S
        )
        ctx = _ctx()
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert resp.source_name == "s2"
        assert ctx.snapshot().hedges == 1

        # 格 3b: 两者皆设 + pre-TTFT 挂起 → 阶段 1 AFTER_S 触发(两阶段各管一段)
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport({"s1": [("hang",), ("ok", "late")], "s2": [("ok", "fast")]})
        mw, _, _ = _harness(
            [s1, s2], transport, hedge_after_s=_STALL_AFTER_S, stall_after_s=_STALL_AFTER_S
        )
        ctx = _ctx()
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert resp.source_name == "s2"
        assert ctx.snapshot().hedges == 1

        # 格 4: 皆无 → 逐字单路(默认关闭回归门的单元格形态)
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport({"s1": [("hang",), ("ok", "solo")]})
        mw, _, _ = _harness([s1, s2], transport)
        ctx = _ctx()
        task = asyncio.ensure_future(mw(_req(ctx=ctx)))
        await transport.entered["s1"].wait()
        transport.release["s1"].set()
        resp = await asyncio.wait_for(task, timeout=5)
        assert resp.content == "solo" and resp.source_name == "s1"
        stats = ctx.snapshot()
        assert stats.hedges == 0 and stats.attempts == 1
        assert transport.ce_events == [None]  # 未配置 = 不观测增量 chunk

    async def test_stall_only_non_stream_degrades_to_single_leg(self):
        """仅 stall 已设 + `stream=False` → 逐字单路(中段沉默无非流式概念,入口门)。"""
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport({"s1": [("hang",), ("ok", "x")], "s2": [("ok", "other")]})
        mw, _, _ = _harness([s1, s2], transport, stall_after_s=_STALL_AFTER_S)
        ctx = _ctx()
        task = asyncio.ensure_future(mw(_req(stream=False, ctx=ctx)))
        await transport.entered["s1"].wait()
        await asyncio.sleep(4 * _STALL_AFTER_S)  # 阈值窗已过,仍不得有第二请求
        assert transport.calls == ["s1"]
        transport.release["s1"].set()
        resp = await asyncio.wait_for(task, timeout=5)
        assert resp.source_name == "s1"
        stats = ctx.snapshot()
        assert stats.hedges == 0 and stats.attempts == 1
        assert transport.ce_events == [None]  # 非流式不建 content_event


class TestStallHedgeWatchdogAndFailure:
    """与看门狗的交互(设计 §5.2): 主路被掐死 ≠ 本轮失败,副路已在跑则等副路。"""

    async def test_watchdog_kills_primary_hedge_leg_survives(self):
        """主路首 chunk 后被看门狗掐死(Transient)、副路在途 → 副路结果裁定。

        attempt 级记账照常(N9): 主路源熔断失败计数 +1、遥测主路 error 行
        (TransientError 对象)+ 副路正常行;"先了结的是失败则等另一路"逐字复用。
        """
        emitter = RecordingEmitter()
        gate = InMemoryGate(config=_BREAKER)
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {
                "s1": [
                    ("chunk",),
                    (
                        "fail_after",
                        0.12,
                        lambda: TransientError("inter_token timeout", source_name="s1"),
                    ),
                ],
                "s2": [("ok_after", 0.2, "rescued")],
            }
        )
        mw, _, _ = _harness(
            [s1, s2], transport, stall_after_s=_STALL_AFTER_S, emitter=emitter, gate=gate
        )
        ctx = _ctx()
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert resp.content == "rescued" and resp.source_name == "s2"
        stats = ctx.snapshot()
        assert stats.hedges == 1 and stats.hedge_won is True
        # 掐死是 attempt 级失败: 熔断失败样本 +1(对冲输家取消不喂,此处是喂的——N9 不对称)
        gate_s1 = gate._gates["s1"]
        assert gate_s1.a0 + gate_s1.a1 == 1
        # 遥测: 主路 error 行携 TransientError 对象(掐死),副路正常行
        primary_row = next(r for r in emitter.rows if r["source"] == "s1")
        hedge_row = next(r for r in emitter.rows if r["source"] == "s2")
        assert isinstance(primary_row["error"], TransientError)
        assert hedge_row["error"] is None

    async def test_both_fail_counts_budget_once(self):
        """主路看门狗掐死 + 副路 Transient: 两败只计一次预算(_combine_failures 复用)。

        max_attempts=2 时恰进第二轮: 若两败计两次,第一轮即耗尽,第二轮的两路
        调用根本不会发生。两败轮次无赢家 → generation 不记、stall 快照恒 None。
        """
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {
                "s1": [
                    ("chunk",),
                    ("fail_after", 0.1, lambda: TransientError("wd", source_name="s1")),
                    ("chunk",),
                    ("fail_after", 0.1, lambda: TransientError("wd", source_name="s1")),
                ],
                "s2": [
                    ("fail_after", 0.12, lambda: TransientError("h", source_name="s2")),
                    ("fail_after", 0.12, lambda: TransientError("h", source_name="s2")),
                ],
            }
        )
        mw, _, _ = _harness([s1, s2], transport, stall_after_s=_STALL_AFTER_S, max_attempts=2)
        ctx = _ctx()
        with pytest.raises(AllSourcesExhausted) as ei:
            await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert ei.value.reason == "retry_exhausted"
        assert transport.calls == ["s1", "s2", "s1", "s2"]  # 两败计一次才有第二轮
        stats = ctx.snapshot()
        assert stats.hedges == 2  # 两败轮次路数照登(实际并发发出即计)
        assert stats.hedge_won is False
        assert stats.generation_ms == 0  # 无赢家: 裸生成时间无归属不记
        assert stats.stall_ms is None and stats.stall_count is None


class TestStallHedgeCancellation:
    """取消穿透(计划取消铁律): stall 等待中外部取消,waiter 收口无残留。"""

    async def test_external_cancel_during_stall_wait_penetrates(self):
        """阶段 2 循环内外部 `task.cancel()`: CancelledError 穿透、无残留任务。

        阈值取 5.0s 宽窗: 取消确定落在 stall 等待内(阶段 2 循环),不依赖时序。
        对冲尚未发出(宽窗未到期)→ 只有原路被取消,permit 释放。
        """
        before = asyncio.all_tasks()
        emitter = RecordingEmitter()
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {"s1": [("chunk",), ("hang",), ("ok", "x")], "s2": [("ok", "y")]}
        )
        mw, limiter, _ = _harness([s1, s2], transport, stall_after_s=5.0, emitter=emitter)
        task = asyncio.ensure_future(mw(_req(ctx=_ctx())))
        await transport.entered["s1"].wait()  # 原路已交付首 chunk 并挂起: 编排处于阶段 2
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert transport.calls == ["s1"]  # 宽窗未到期,对冲从未发出
        assert {r["source"]: r["error"] for r in emitter.rows} == {"s1": "cancelled"}
        assert (await limiter.source_stats("s1")).inflight == 0
        assert asyncio.all_tasks() == before  # waiter/原路全收口,无残留


class TestStallHedgeEventDiscipline:
    """Event 并发纪律(N3③): 对冲路 content_event 恒 None(v1 不梯次)。"""

    async def test_hedge_leg_content_event_always_none(self):
        """原路(流式 + stall 已设)收真事件,对冲路恒 None——逐次身份记录断言。"""
        s1, s2 = _src("s1"), _src("s2")
        transport = StallHedgeTransport(
            {"s1": [("chunk",), ("hang",), ("ok", "late")], "s2": [("ok", "win")]}
        )
        mw, _, _ = _harness([s1, s2], transport, stall_after_s=_STALL_AFTER_S)
        ctx = _ctx()
        resp = await asyncio.wait_for(mw(_req(ctx=ctx)), timeout=5)
        assert resp.source_name == "s2"
        assert transport.calls == ["s1", "s2"]
        # 原路携带首 token + content 双观测;对冲路两事件恒 None(N3③,同 H5 首 token 约定)
        assert [e is not None for e in transport.ft_events] == [True, False]
        assert [e is not None for e in transport.ce_events] == [True, False]