feat: add retry middleware with per-attempt governance orchestration
This commit is contained in:
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user