From c3d5079d39abe6ec697bc8293973cb02db27a5a3 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Mon, 20 Jul 2026 07:05:42 -0400 Subject: [PATCH] feat: add retry middleware with per-attempt governance orchestration --- src/polygateway/middleware/base.py | 25 ++ src/polygateway/middleware/breaker.py | 58 +++++ src/polygateway/middleware/ratelimit.py | 46 ++++ src/polygateway/middleware/retry.py | 295 ++++++++++++++++++++++++ tests/unit/test_retry.py | 281 ++++++++++++++++++++++ 5 files changed, 705 insertions(+) create mode 100644 src/polygateway/middleware/base.py create mode 100644 src/polygateway/middleware/breaker.py create mode 100644 src/polygateway/middleware/ratelimit.py create mode 100644 src/polygateway/middleware/retry.py create mode 100644 tests/unit/test_retry.py diff --git a/src/polygateway/middleware/base.py b/src/polygateway/middleware/base.py new file mode 100644 index 0000000..14a9e4a --- /dev/null +++ b/src/polygateway/middleware/base.py @@ -0,0 +1,25 @@ +"""洋葱组装(D1): 把中间件序列外→内绑定到终端调用上。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Sequence + + from polygateway.ports import CallNext, Middleware + + +def compose(middlewares: Sequence[Middleware], terminal: CallNext) -> CallNext: + """外→内绑定: compose([A, B], t) 的调用序为 A → B → t。""" + handler = terminal + for mw in reversed(middlewares): + handler = _bind(mw, handler) + return handler + + +def _bind(mw: Middleware, nxt: CallNext) -> CallNext: + async def call(request): + return await mw(request, nxt) + + return call diff --git a/src/polygateway/middleware/breaker.py b/src/polygateway/middleware/breaker.py new file mode 100644 index 0000000..2061aba --- /dev/null +++ b/src/polygateway/middleware/breaker.py @@ -0,0 +1,58 @@ +"""熔断门封装: 后端故障必须报错而非放行(降级方向铁律,与 QuotaGate 同构)。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from polygateway.errors import GovernanceBackendError + +if TYPE_CHECKING: + from polygateway.ports import GateDecision, GateUpdate, ProviderGate + from polygateway.types import SourceConfig + + +class BreakerGate: + """RetryMW 面向熔断后端的唯一入口;包装一切后端异常。""" + + def __init__(self, gate: ProviderGate) -> None: + self._gate = gate + + async def try_enter(self, source: SourceConfig, owner: str) -> GateDecision: + try: + return await self._gate.try_enter(source.name, owner) + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"熔断后端故障(try_enter): {exc}") from exc + + async def record_success(self, entry: GateDecision) -> GateUpdate: + try: + return await self._gate.record_success(entry) + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"熔断后端故障(record_success): {exc}") from exc + + async def record_failure(self, entry: GateDecision, reason: str, force_open: bool) -> GateUpdate: + try: + return await self._gate.record_failure(entry, reason, force_open) + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"熔断后端故障(record_failure): {exc}") from exc + + async def release_probe(self, entry: GateDecision) -> GateUpdate: + try: + return await self._gate.release_probe(entry) + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"熔断后端故障(release_probe): {exc}") from exc + + async def retry_after_s(self, sources: tuple[str, ...]) -> float: + try: + return await self._gate.retry_after_s(sources) + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"熔断后端故障(retry_after_s): {exc}") from exc diff --git a/src/polygateway/middleware/ratelimit.py b/src/polygateway/middleware/ratelimit.py new file mode 100644 index 0000000..6035dcf --- /dev/null +++ b/src/polygateway/middleware/ratelimit.py @@ -0,0 +1,46 @@ +"""限流准入封装: 后端故障必须报错而非放行(降级方向铁律)。 + +非降级设计说明: 缓存/遥测挂 → 静默;限流挂 → 若放行,多 worker 会同时 +失去闸门直接击穿网关,故一律 GovernanceBackendError 上抛。 +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from polygateway.errors import GovernanceBackendError + +if TYPE_CHECKING: + from polygateway.ports import Permit, RateLimiter + from polygateway.types import SourceConfig, SourceStats + + +class QuotaGate: + """RetryMW 面向限流后端的唯一入口;包装一切后端异常。""" + + def __init__(self, limiter: RateLimiter) -> None: + self._limiter = limiter + + async def try_acquire(self, source: SourceConfig) -> Permit | None: + try: + return await self._limiter.try_acquire(source.name, source.est_tokens) + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"限流后端故障(try_acquire): {exc}") from exc + + async def stats(self, source: SourceConfig) -> SourceStats: + try: + return await self._limiter.source_stats(source.name) + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"限流后端故障(source_stats): {exc}") from exc + + async def mark_progress(self) -> None: + try: + await self._limiter.mark_progress() + except GovernanceBackendError: + raise + except Exception as exc: + raise GovernanceBackendError(f"限流后端故障(mark_progress): {exc}") from exc diff --git a/src/polygateway/middleware/retry.py b/src/polygateway/middleware/retry.py new file mode 100644 index 0000000..c119d97 --- /dev/null +++ b/src/polygateway/middleware/retry.py @@ -0,0 +1,295 @@ +"""RetryMW: 重试循环拥有"尝试"的全部编排(D13 自研;蓝本 CHS governance.py:107-268)。 + +每次尝试 = 选源(跳过冷却源)→ 该源熔断门 → 限流 permit → transport; +换源、退避、逐次遥测、熔断写回、permit 结算全部在循环内。作为洋葱的 +终端(compose 的 terminal)被调用: `await retry_mw(request)`。 + +例外说明(P7): 本模块 import httpx 仅用于失败原因归类(CHS 同款判定), +httpx 是库的核心依赖而非实现层内部件,不违反"middleware 只依赖端口"。 +""" + +from __future__ import annotations + +import asyncio +import random +import time +import uuid +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import httpx +from loguru import logger + +from polygateway.errors import ( + AllSourcesExhausted, + CircuitOpenError, + PolyGatewayError, + RequestRejectedError, + ResultInvalidError, + SourceDeadError, + TransientError, +) +from polygateway.middleware.breaker import BreakerGate +from polygateway.middleware.ratelimit import QuotaGate +from polygateway.sources import SourceCooldownMemo +from polygateway.streaming import StreamLivenessTimeout +from polygateway.types import LLMResponse + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from polygateway.ports import ( + GateDecision, + Permit, + ProviderGate, + RateLimiter, + SourceSelector, + Transport, + ) + from polygateway.types import ( + BackpressurePolicy, + ChatRequest, + RetryPolicy, + SourceConfig, + TransportResult, + ) + + +def _failure_reason(exc: PolyGatewayError) -> str: + """失败原因归类(CHS governance.py:169 同款)。""" + if isinstance(exc, SourceDeadError): + return "source_dead" + if exc.status_code == 429: + return "rate_limited" + if isinstance(exc.__cause__, httpx.TimeoutException | StreamLivenessTimeout): + return "timeout" + return "network_error" + + +@dataclass(frozen=True) +class _Failed: + """一次可重试失败的结果(SourceDead 立即换源,Transient 退避)。""" + + exc: PolyGatewayError + immediate: bool + + +class RetryMW: + """尝试编排器;时钟/睡眠/随机全部注入,纯确定性可测(P6)。""" + + def __init__( + self, + *, + scope: str, + sources: list[SourceConfig], + selector: SourceSelector, + limiter: RateLimiter, + gate: ProviderGate, + transport: Transport, + retry: RetryPolicy, + backpressure: BackpressurePolicy, + quota_full: str = "wait", + cooldown_memo: SourceCooldownMemo | None = None, + emitter: object | None = None, + now: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + rng: Callable[[], float] = random.random, + ) -> None: + if quota_full not in ("wait", "fail_fast"): + raise ValueError(f"quota_full 必须是 wait|fail_fast: {quota_full!r}") + self._scope = scope + self._sources = list(sources) + self._selector = selector + self._quota = QuotaGate(limiter) + self._breaker = BreakerGate(gate) + self._transport = transport + self._retry = retry + self._bp = backpressure + self._quota_full = quota_full + self._memo = cooldown_memo or SourceCooldownMemo(now=now) + self._emitter = emitter + self._now = now + self._sleep = sleep + self._rng = rng + + async def __call__(self, request: ChatRequest) -> LLMResponse: + """执行治理调用;scope 级失败按 §6.1 携结构化字段上抛。""" + if not self._sources: + raise AllSourcesExhausted(scope=self._scope, reason="no_sources", retry_after_s=0.0) + fails = 0 + reasons: dict[str, str] = {} + while True: + picked, gate_rejections = await self._pick_runnable(reasons) + if picked is None: + await self._on_no_runnable(gate_rejections, reasons) + continue + outcome = await self._attempt(request, *picked, reasons) + if isinstance(outcome, LLMResponse): + return outcome + fails += 1 + if fails >= self._retry.max_attempts: + raise AllSourcesExhausted( + scope=self._scope, reason="retry_exhausted", + retry_after_s=self._retry.backoff_base_s, + per_source_reasons=reasons, + ) from outcome.exc + if not outcome.immediate: + await self._sleep(self._backoff_delay(fails, outcome.exc)) + + # —— 选源与准入(CHS _pick_runnable 120-167)—— + + async def _pick_runnable( + self, reasons: dict[str, str] + ) -> tuple[tuple[SourceConfig, Permit, GateDecision] | None, int]: + stats = {s.name: await self._quota.stats(s) for s in self._sources} + gate_rejections = 0 + for cand in self._selector.order(self._sources, stats): + if self._memo.active(cand.name): + # 冷却备忘跳过也计入拒绝数,保住 circuit_open 判据(CHS 同款) + gate_rejections += 1 + reasons[cand.name] = "cooldown" + continue + permit = await self._quota.try_acquire(cand) + if permit is None: + reasons.setdefault(cand.name, "rate_limited") + continue + try: + entry = await self._breaker.try_enter(cand, uuid.uuid4().hex) + except BaseException: + await self._settle_and_release(permit, 0) + raise + if entry.allowed: + return (cand, permit, entry), gate_rejections + gate_rejections += 1 + reasons[cand.name] = "circuit_open" + # 开路源本地记冷却,避免每轮白烧 RPM 探测(CHS governance.py:107) + self._memo.set_until(cand.name, self._now() + entry.retry_after_s) + await self._settle_and_release(permit, 0) + return None, gate_rejections + + async def _on_no_runnable(self, gate_rejections: int, reasons: dict[str, str]) -> None: + if gate_rejections == len(self._sources): + names = tuple(s.name for s in self._sources) + raise CircuitOpenError( + scope=self._scope, + retry_after_s=await self._breaker.retry_after_s(names), + per_source_reasons=reasons, + ) + if self._quota_full == "fail_fast": + raise AllSourcesExhausted( + scope=self._scope, reason="quota_exhausted", + retry_after_s=self._bp.poll_interval_s, per_source_reasons=reasons, + ) + await self._sleep(self._bp.poll_interval_s) + + # —— 单次尝试(CHS run 200-268)—— + + async def _attempt( + self, + request: ChatRequest, + source: SourceConfig, + permit: Permit, + entry: GateDecision, + reasons: dict[str, str], + ) -> LLMResponse | _Failed: + call_id = str(uuid.uuid4()) + started = self._now() + actual = 0 + try: + result = await self._transport.complete( + messages=request.messages, source=source, + stream=request.stream, overlay=request.overlay, call_id=call_id, + ) + actual = result.prompt_tokens + result.completion_tokens + await self._breaker.record_success(entry) + await self._quota.mark_progress() + response = self._build_response(source, result, call_id, started) + await self._emit(request, source, call_id, started, response=response) + return response + except RequestRejectedError as exc: + await self._on_rejected(exc, source, entry) + await self._emit(request, source, call_id, started, error=exc) + raise + except ResultInvalidError as exc: + # 坏结果 ≠ 坏服务: 熔断记成功,异常上抛消耗业务失败预算(§6.3) + await self._breaker.record_success(entry) + await self._emit(request, source, call_id, started, error=exc) + raise + except asyncio.CancelledError: + if entry.is_probe: + await self._breaker.release_probe(entry) + await self._emit(request, source, call_id, started, error="cancelled") + raise + except (SourceDeadError, TransientError) as exc: + dead = isinstance(exc, SourceDeadError) + reason = _failure_reason(exc) + reasons[source.name] = reason + await self._breaker.record_failure(entry, reason, dead) + if not dead: + actual = source.est_tokens # 保守: 失败请求可能已被网关计费(CHS 同款) + await self._emit(request, source, call_id, started, error=exc) + return _Failed(exc, immediate=dead) + finally: + await self._settle_and_release(permit, actual) + + async def _on_rejected( + self, exc: RequestRejectedError, source: SourceConfig, entry: GateDecision + ) -> None: + provider_responded = exc.source_name == source.name and exc.status_code is not None + if provider_responded: + await self._breaker.record_success(entry) # 网关健康地拒了坏请求 + elif entry.is_probe: + await self._breaker.release_probe(entry) + + # —— 辅助 —— + + def _backoff_delay(self, fails: int, exc: PolyGatewayError) -> float: + """指数退避+jitter,与 Retry-After 提示取大(ARCH §7.2;VT jitter 系数)。""" + base = min(self._retry.backoff_base_s * (2 ** (fails - 1)), self._retry.backoff_max_s) + delay = base * (0.5 + self._rng()) + retry_after = getattr(exc, "retry_after_s", None) or 0.0 + return max(delay, retry_after) + + def _build_response( + self, source: SourceConfig, result: TransportResult, call_id: str, started: float + ) -> LLMResponse: + return LLMResponse( + content=result.content, thinking=result.thinking, + model=source.model, provider=source.provider, + prompt_tokens=result.prompt_tokens, completion_tokens=result.completion_tokens, + latency_ms=int((self._now() - started) * 1000), + ttft_ms=result.ttft_ms, max_inter_token_ms=result.max_inter_token_ms, + cache_hit=False, call_id=call_id, + source_name=source.name, cost=None, usage_source=result.usage_source, + ) + + async def _settle_and_release(self, permit: Permit, actual: int) -> None: + """finally 专用: settle 后必 release;失败降级 warning,绝不掩盖主异常/取消。""" + try: + try: + await permit.settle(actual) + finally: + await permit.release() + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("permit 结算/释放失败(不掩盖主异常): {}", exc) + + async def _emit( + self, request: ChatRequest, source: SourceConfig, call_id: str, started: float, + *, response: LLMResponse | None = None, error: object | None = None, + ) -> None: + """逐次遥测(经注入的单一 Emitter);遥测失败不得影响调用(铁律)。""" + if self._emitter is None: + return + try: + await self._emitter.emit_attempt( + request=request, source=source, call_id=call_id, + latency_ms=int((self._now() - started) * 1000), + response=response, + error=None if error is None else str(error), + ) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("逐次遥测记录失败(降级不冒泡): {}", exc) diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py new file mode 100644 index 0000000..677a1c0 --- /dev/null +++ b/tests/unit/test_retry.py @@ -0,0 +1,281 @@ +"""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.ports import GateState +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 = dict( + 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): + 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=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_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