mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
1845 lines
72 KiB
Python
1845 lines
72 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import re
|
||
from collections.abc import Awaitable, Callable
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timezone
|
||
from typing import Any, AsyncIterator
|
||
|
||
from sqlalchemy import update
|
||
from sqlalchemy.orm import Session
|
||
|
||
from src.core.logger import logger
|
||
from src.models.database import RequestCandidate
|
||
from src.services.orchestration.error_classifier import ErrorAction, ErrorClassifier
|
||
from src.services.request.candidate import RequestCandidateService
|
||
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
|
||
from src.services.task.core.exceptions import StreamProbeError
|
||
from src.services.task.core.protocol import AttemptFunc, AttemptKind, AttemptResult
|
||
from src.services.task.core.schema import ExecutionResult
|
||
|
||
from .policy import FailoverAction, RetryMode, RetryPolicy, SkipPolicy
|
||
from .recorder import CandidateRecorder
|
||
from .schema import CandidateKey
|
||
|
||
_DISCONNECT_EXCEPTION_NAMES = frozenset({"ClientDisconnectedException"})
|
||
|
||
|
||
def _is_client_disconnected(exc: Exception) -> bool:
|
||
"""检查异常(或其 cause)是否为客户端断连，避免循环 import。"""
|
||
for obj in (exc, getattr(exc, "cause", None)):
|
||
if obj is not None and type(obj).__name__ in _DISCONNECT_EXCEPTION_NAMES:
|
||
return True
|
||
return False
|
||
|
||
|
||
_SENSITIVE_PATTERN = re.compile(
|
||
r"(api[_-]?key|token|bearer|authorization)[=:\s]+\S+",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class AttemptErrorOutcome:
|
||
"""_handle_attempt_error 的返回结果"""
|
||
|
||
action: FailoverAction
|
||
last_status_code: int | None
|
||
max_retries: int
|
||
stop_result: ExecutionResult | None = None
|
||
|
||
|
||
class FailoverEngine:
|
||
"""
|
||
FailoverEngine executes candidate attempts under policies.
|
||
|
||
Phase3 core: unified failover loop used by TaskService.
|
||
"""
|
||
|
||
# Hard constraint: streaming first chunk probe timeout
|
||
STREAM_FIRST_CHUNK_TIMEOUT_SECONDS: int = 30
|
||
RETRY_BACKOFF_EVERY_FAILURES: int = 10
|
||
RETRY_ROTATE_CLIENT_EVERY_FAILURES: int = 40
|
||
RETRY_BACKOFF_BASE_SECONDS: float = 0.025
|
||
RETRY_BACKOFF_MAX_SECONDS: float = 0.15
|
||
STREAM_CAPACITY_BACKOFF_SECONDS: float = 0.3
|
||
|
||
def __init__(
|
||
self,
|
||
db: Session,
|
||
*,
|
||
error_classifier: ErrorClassifier | None = None,
|
||
recorder: CandidateRecorder | None = None,
|
||
) -> None:
|
||
self.db = db
|
||
self._error_classifier = error_classifier or ErrorClassifier(db=db)
|
||
self._recorder = recorder or CandidateRecorder(db)
|
||
|
||
async def _db_op(self, func: Callable[[], Any]) -> Any:
|
||
"""将同步 DB 操作放到线程池执行，避免阻塞 asyncio 事件循环。
|
||
|
||
当事件循环被同步 db.commit() / db.execute() 阻塞时，
|
||
tunnel relay PING 心跳无法发送，导致 worker idle timeout 断连，整个服务不可用。
|
||
"""
|
||
return await asyncio.to_thread(func)
|
||
|
||
@staticmethod
|
||
def _collect_error_messages(error: Exception | None) -> str:
|
||
if error is None:
|
||
return ""
|
||
|
||
parts: list[str] = []
|
||
for item in (
|
||
getattr(error, "message", None),
|
||
getattr(error, "upstream_response", None),
|
||
str(error),
|
||
):
|
||
if isinstance(item, str) and item.strip():
|
||
parts.append(item.strip())
|
||
|
||
cause = getattr(error, "cause", None)
|
||
if cause is not None and cause is not error:
|
||
for item in (
|
||
getattr(cause, "message", None),
|
||
getattr(cause, "upstream_response", None),
|
||
str(cause),
|
||
):
|
||
if isinstance(item, str) and item.strip():
|
||
parts.append(item.strip())
|
||
|
||
return " | ".join(parts)
|
||
|
||
@classmethod
|
||
def _is_stream_capacity_error(cls, error: Exception | None) -> bool:
|
||
lowered = cls._collect_error_messages(error).lower()
|
||
return (
|
||
"max outbound streams" in lowered
|
||
or "too many concurrent streams" in lowered
|
||
or "max concurrent streams" in lowered
|
||
)
|
||
|
||
@classmethod
|
||
def _compute_retry_backoff_seconds(
|
||
cls,
|
||
*,
|
||
consecutive_failures: int,
|
||
error: Exception | None,
|
||
) -> float:
|
||
if consecutive_failures <= 0:
|
||
return 0.0
|
||
if cls._is_stream_capacity_error(error):
|
||
return cls.STREAM_CAPACITY_BACKOFF_SECONDS
|
||
if consecutive_failures % cls.RETRY_BACKOFF_EVERY_FAILURES != 0:
|
||
return 0.0
|
||
step = max(1, consecutive_failures // cls.RETRY_BACKOFF_EVERY_FAILURES)
|
||
return min(cls.RETRY_BACKOFF_BASE_SECONDS * step, cls.RETRY_BACKOFF_MAX_SECONDS)
|
||
|
||
@classmethod
|
||
def _should_rotate_upstream_client(
|
||
cls,
|
||
*,
|
||
consecutive_failures: int,
|
||
error: Exception | None,
|
||
) -> bool:
|
||
if cls._is_stream_capacity_error(error):
|
||
return True
|
||
return (
|
||
consecutive_failures >= cls.RETRY_ROTATE_CLIENT_EVERY_FAILURES
|
||
and consecutive_failures % cls.RETRY_ROTATE_CLIENT_EVERY_FAILURES == 0
|
||
)
|
||
|
||
async def _rotate_upstream_client(self, candidate: ProviderCandidate) -> bool:
|
||
from src.clients.http_client import HTTPClientPool
|
||
from src.services.proxy_node.resolver import (
|
||
resolve_delegate_config_async,
|
||
resolve_effective_proxy,
|
||
)
|
||
|
||
effective_proxy = resolve_effective_proxy(
|
||
getattr(candidate.provider, "proxy", None),
|
||
getattr(candidate.key, "proxy", None),
|
||
)
|
||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||
return await HTTPClientPool.reset_upstream_client(
|
||
delegate_cfg, proxy_config=effective_proxy
|
||
)
|
||
|
||
async def _apply_retry_pacing(
|
||
self,
|
||
*,
|
||
candidate: ProviderCandidate,
|
||
consecutive_failures: int,
|
||
error: Exception | None,
|
||
request_id: str | None,
|
||
) -> None:
|
||
should_rotate = self._should_rotate_upstream_client(
|
||
consecutive_failures=consecutive_failures,
|
||
error=error,
|
||
)
|
||
if should_rotate:
|
||
rotated = await self._rotate_upstream_client(candidate)
|
||
if rotated:
|
||
logger.warning(
|
||
" [{}] 连续失败 {} 次，已重建上游客户端复用",
|
||
request_id,
|
||
consecutive_failures,
|
||
)
|
||
|
||
backoff_seconds = self._compute_retry_backoff_seconds(
|
||
consecutive_failures=consecutive_failures,
|
||
error=error,
|
||
)
|
||
if backoff_seconds > 0:
|
||
logger.warning(
|
||
" [{}] 连续失败 {} 次，退避 {:.0f}ms 后继续尝试",
|
||
request_id,
|
||
consecutive_failures,
|
||
backoff_seconds * 1000,
|
||
)
|
||
await asyncio.sleep(backoff_seconds)
|
||
|
||
async def _check_cancellation(
|
||
self,
|
||
is_cancelled: Callable[[], Awaitable[bool]] | None,
|
||
) -> bool:
|
||
if is_cancelled is None:
|
||
return False
|
||
try:
|
||
return bool(await is_cancelled())
|
||
except Exception:
|
||
return False
|
||
|
||
async def _mark_remaining_cancelled(
|
||
self,
|
||
*,
|
||
candidate_record_map: dict[tuple[int, int], str] | None,
|
||
candidates: list[ProviderCandidate],
|
||
from_candidate_idx: int,
|
||
from_retry_idx: int,
|
||
retry_policy: RetryPolicy,
|
||
) -> None:
|
||
if not candidate_record_map:
|
||
return
|
||
|
||
# Pre-compute outside thread (reads ProviderCandidate attrs that may not be thread-safe)
|
||
record_ids: list[str] = []
|
||
for candidate_idx, cand in enumerate(candidates):
|
||
if candidate_idx < from_candidate_idx:
|
||
continue
|
||
max_retries = self._get_max_retries(cand, retry_policy)
|
||
for retry_idx in range(max_retries):
|
||
if candidate_idx == from_candidate_idx and retry_idx < from_retry_idx:
|
||
continue
|
||
record_id = candidate_record_map.get((candidate_idx, retry_idx))
|
||
if record_id:
|
||
record_ids.append(record_id)
|
||
|
||
if not record_ids:
|
||
return
|
||
|
||
def _do() -> None:
|
||
now = datetime.now(timezone.utc)
|
||
for rid in record_ids:
|
||
self.db.execute(
|
||
update(RequestCandidate)
|
||
.where(RequestCandidate.id == rid)
|
||
.where(RequestCandidate.status.in_(["available", "pending"]))
|
||
.values(
|
||
status="cancelled",
|
||
status_code=499,
|
||
error_message="cancelled_by_client",
|
||
finished_at=now,
|
||
)
|
||
)
|
||
self.db.commit()
|
||
|
||
await self._db_op(_do)
|
||
|
||
def _append_cancelled_fallback_candidate_keys(
|
||
self,
|
||
*,
|
||
fallback: list[CandidateKey],
|
||
candidates: list[ProviderCandidate],
|
||
from_candidate_idx: int,
|
||
from_retry_idx: int,
|
||
retry_policy: RetryPolicy,
|
||
) -> None:
|
||
existing = {(item.candidate_index, item.retry_index) for item in fallback}
|
||
for candidate_idx, cand in enumerate(candidates):
|
||
if candidate_idx < from_candidate_idx:
|
||
continue
|
||
max_retries = self._get_max_retries(cand, retry_policy)
|
||
for retry_idx in range(max_retries):
|
||
if candidate_idx == from_candidate_idx and retry_idx < from_retry_idx:
|
||
continue
|
||
key = (candidate_idx, retry_idx)
|
||
if key in existing:
|
||
continue
|
||
original_key = getattr(cand, "key", None)
|
||
original_pool_key_index = getattr(cand, "_pool_key_index", 0)
|
||
if isinstance(cand, PoolCandidate) and cand.pool_keys:
|
||
retry_slots_per_key = self._get_pool_key_max_retries(cand, retry_policy)
|
||
pool_key_index = min(retry_idx // retry_slots_per_key, len(cand.pool_keys) - 1)
|
||
cand.key = cand.pool_keys[pool_key_index]
|
||
cand._pool_key_index = pool_key_index
|
||
fallback.append(
|
||
self._make_candidate_key(
|
||
candidate=cand,
|
||
candidate_index=candidate_idx,
|
||
retry_index=retry_idx,
|
||
status="cancelled",
|
||
error_message="cancelled_by_client",
|
||
status_code=499,
|
||
)
|
||
)
|
||
if isinstance(cand, PoolCandidate):
|
||
cand.key = original_key
|
||
cand._pool_key_index = original_pool_key_index
|
||
existing.add(key)
|
||
|
||
async def _maybe_cancel_execution(
|
||
self,
|
||
*,
|
||
is_cancelled: Callable[[], Awaitable[bool]] | None,
|
||
candidate_record_map: dict[tuple[int, int], str] | None,
|
||
candidate_keys_fallback: list[CandidateKey],
|
||
candidates: list[ProviderCandidate],
|
||
from_candidate_idx: int,
|
||
from_retry_idx: int,
|
||
retry_policy: RetryPolicy,
|
||
request_id: str | None,
|
||
attempt_count: int,
|
||
) -> ExecutionResult | None:
|
||
if not await self._check_cancellation(is_cancelled):
|
||
return None
|
||
|
||
logger.info(
|
||
"[FailoverEngine] Request cancelled by client at candidate_index={}, retry_index={}",
|
||
from_candidate_idx,
|
||
from_retry_idx,
|
||
)
|
||
await self._mark_remaining_cancelled(
|
||
candidate_record_map=candidate_record_map,
|
||
candidates=candidates,
|
||
from_candidate_idx=from_candidate_idx,
|
||
from_retry_idx=from_retry_idx,
|
||
retry_policy=retry_policy,
|
||
)
|
||
self._append_cancelled_fallback_candidate_keys(
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=from_candidate_idx,
|
||
from_retry_idx=from_retry_idx,
|
||
retry_policy=retry_policy,
|
||
)
|
||
return ExecutionResult(
|
||
success=False,
|
||
error_type="cancelled",
|
||
error_message="cancelled_by_client",
|
||
last_status_code=499,
|
||
candidate_keys=self._get_candidate_keys(
|
||
request_id=request_id,
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
),
|
||
attempt_count=attempt_count,
|
||
)
|
||
|
||
async def _build_disconnected_result(
|
||
self,
|
||
*,
|
||
candidate_record_map: dict[tuple[int, int], str] | None,
|
||
candidate_keys_fallback: list[CandidateKey],
|
||
candidates: list[ProviderCandidate],
|
||
from_candidate_idx: int,
|
||
from_retry_idx: int,
|
||
retry_policy: RetryPolicy,
|
||
request_id: str | None,
|
||
attempt_count: int,
|
||
) -> ExecutionResult:
|
||
"""客户端已断连，立即终止故障转移并返回结果。"""
|
||
await self._mark_remaining_cancelled(
|
||
candidate_record_map=candidate_record_map,
|
||
candidates=candidates,
|
||
from_candidate_idx=from_candidate_idx,
|
||
from_retry_idx=from_retry_idx,
|
||
retry_policy=retry_policy,
|
||
)
|
||
self._append_cancelled_fallback_candidate_keys(
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=from_candidate_idx,
|
||
from_retry_idx=from_retry_idx,
|
||
retry_policy=retry_policy,
|
||
)
|
||
return ExecutionResult(
|
||
success=False,
|
||
error_type="ClientDisconnected",
|
||
error_message="client_disconnected",
|
||
last_status_code=499,
|
||
candidate_keys=self._get_candidate_keys(
|
||
request_id=request_id,
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
),
|
||
attempt_count=attempt_count,
|
||
)
|
||
|
||
async def execute(
|
||
self,
|
||
*,
|
||
candidates: list[ProviderCandidate],
|
||
attempt_func: AttemptFunc,
|
||
retry_policy: RetryPolicy,
|
||
skip_policy: SkipPolicy,
|
||
request_id: str | None = None,
|
||
user_id: str | None = None,
|
||
api_key_id: str | None = None,
|
||
username: str | None = None,
|
||
api_key_name: str | None = None,
|
||
candidate_record_map: dict[tuple[int, int], str] | None = None,
|
||
max_candidates: int | None = None,
|
||
max_attempts: int | None = None,
|
||
execution_error_handler: (
|
||
Callable[
|
||
...,
|
||
Awaitable[tuple[FailoverAction, int | None]],
|
||
]
|
||
| None
|
||
) = None,
|
||
is_cancelled: Callable[[], Awaitable[bool]] | None = None,
|
||
) -> ExecutionResult:
|
||
"""
|
||
Execute candidate traversal + retry + failover.
|
||
|
||
Notes:
|
||
- For PRE_EXPAND: `candidate_record_map` should be provided (created by CandidateResolver).
|
||
- For ON_DEMAND/DISABLED: records are created when used (and on skip, best-effort).
|
||
"""
|
||
candidate_keys_fallback: list[CandidateKey] = []
|
||
|
||
if max_candidates is not None and max_candidates > 0:
|
||
candidates = candidates[:max_candidates]
|
||
|
||
attempt_count = 0
|
||
consecutive_failures = 0
|
||
last_status_code: int | None = None
|
||
|
||
# For logging / dispatcher parity only; callers may pass an exact value.
|
||
if max_attempts is None:
|
||
computed = 0
|
||
for cand in candidates:
|
||
should_skip, _ = self._should_skip(cand, skip_policy)
|
||
if should_skip:
|
||
continue
|
||
computed += self._get_max_retries(cand, retry_policy)
|
||
max_attempts = computed
|
||
|
||
for candidate_index, candidate in enumerate(candidates):
|
||
cancelled_result = await self._maybe_cancel_execution(
|
||
is_cancelled=is_cancelled,
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=candidate_index,
|
||
from_retry_idx=0,
|
||
retry_policy=retry_policy,
|
||
request_id=request_id,
|
||
attempt_count=attempt_count,
|
||
)
|
||
if cancelled_result is not None:
|
||
return cancelled_result
|
||
|
||
should_skip, skip_reason = self._should_skip(candidate, skip_policy)
|
||
if should_skip:
|
||
# PRE_EXPAND: mark all retry slots skipped.
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_candidate_skipped(
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_index=candidate_index,
|
||
candidate=candidate,
|
||
retry_policy=retry_policy,
|
||
skip_reason=skip_reason,
|
||
)
|
||
else:
|
||
# ON_DEMAND/DISABLED: create a skipped record for audit (best-effort).
|
||
if request_id:
|
||
await self._create_skipped_record(
|
||
request_id=request_id,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
user_id=user_id,
|
||
api_key_id=api_key_id,
|
||
username=username,
|
||
api_key_name=api_key_name,
|
||
skip_reason=skip_reason,
|
||
)
|
||
candidate_keys_fallback.append(
|
||
self._make_candidate_key(
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=0,
|
||
status="skipped",
|
||
skip_reason=skip_reason,
|
||
)
|
||
)
|
||
continue
|
||
|
||
if isinstance(candidate, PoolCandidate):
|
||
pool_result, attempt_count, consecutive_failures, last_status_code = (
|
||
await self._execute_pool_candidate(
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
attempt_func=attempt_func,
|
||
retry_policy=retry_policy,
|
||
request_id=request_id,
|
||
user_id=user_id,
|
||
api_key_id=api_key_id,
|
||
username=username,
|
||
api_key_name=api_key_name,
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
attempt_count=attempt_count,
|
||
max_attempts=max_attempts,
|
||
execution_error_handler=execution_error_handler,
|
||
consecutive_failures=consecutive_failures,
|
||
is_cancelled=is_cancelled,
|
||
)
|
||
)
|
||
if pool_result is not None:
|
||
return pool_result
|
||
continue
|
||
|
||
max_retries = self._get_max_retries(candidate, retry_policy)
|
||
retry_index = 0
|
||
while retry_index < max_retries:
|
||
cancelled_result = await self._maybe_cancel_execution(
|
||
is_cancelled=is_cancelled,
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=candidate_index,
|
||
from_retry_idx=retry_index,
|
||
retry_policy=retry_policy,
|
||
request_id=request_id,
|
||
attempt_count=attempt_count,
|
||
)
|
||
if cancelled_result is not None:
|
||
return cancelled_result
|
||
|
||
attempt_count += 1
|
||
|
||
# Resolve/create record_id
|
||
record_id = None
|
||
if candidate_record_map:
|
||
record_id = candidate_record_map.get((candidate_index, retry_index))
|
||
if record_id is None:
|
||
# Rectify may extend retries beyond pre-created range; reuse retry 0 record.
|
||
record_id = candidate_record_map.get((candidate_index, 0))
|
||
if record_id is None and request_id and retry_policy.mode != RetryMode.PRE_EXPAND:
|
||
record_id = await self._ensure_record_exists(
|
||
request_id=request_id,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
user_id=user_id,
|
||
api_key_id=api_key_id,
|
||
username=username,
|
||
api_key_name=api_key_name,
|
||
)
|
||
|
||
self._attach_attempt_context(
|
||
candidate, candidate_index, retry_index, record_id, attempt_count, max_attempts
|
||
)
|
||
|
||
# Mark pending + commit BEFORE await
|
||
# (avoid holding DB connections during slow upstream calls)
|
||
await self._mark_pending_and_commit(record_id)
|
||
|
||
try:
|
||
attempt_result = await self._execute_attempt(
|
||
candidate=candidate,
|
||
record_id=record_id,
|
||
attempt_func=attempt_func,
|
||
)
|
||
last_status_code = int(getattr(attempt_result, "http_status", 0) or 0)
|
||
|
||
# PRE_EXPAND: mark unused slots after request ends (success)
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_remaining_slots_unused(
|
||
candidate_record_map=candidate_record_map,
|
||
candidates=candidates,
|
||
success_candidate_idx=candidate_index,
|
||
success_retry_idx=retry_index,
|
||
retry_policy=retry_policy,
|
||
)
|
||
|
||
consecutive_failures = 0
|
||
return ExecutionResult(
|
||
success=True,
|
||
attempt_result=attempt_result,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
provider_id=str(candidate.provider.id),
|
||
provider_name=str(candidate.provider.name),
|
||
endpoint_id=str(candidate.endpoint.id),
|
||
key_id=str(candidate.key.id),
|
||
candidate_keys=self._get_candidate_keys(
|
||
request_id=request_id,
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
),
|
||
attempt_count=attempt_count,
|
||
request_candidate_id=record_id,
|
||
)
|
||
|
||
except StreamProbeError as exc:
|
||
last_status_code = exc.http_status
|
||
await self._record_attempt_failure(record_id, exc, exc.http_status)
|
||
action = FailoverAction.CONTINUE
|
||
consecutive_failures += 1
|
||
await self._apply_retry_pacing(
|
||
candidate=candidate,
|
||
consecutive_failures=consecutive_failures,
|
||
error=exc,
|
||
request_id=request_id,
|
||
)
|
||
|
||
except Exception as exc:
|
||
if _is_client_disconnected(exc):
|
||
await self._record_attempt_failure(record_id, exc, 499)
|
||
return await self._build_disconnected_result(
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=candidate_index,
|
||
from_retry_idx=retry_index + 1,
|
||
retry_policy=retry_policy,
|
||
request_id=request_id,
|
||
attempt_count=attempt_count,
|
||
)
|
||
|
||
outcome = await self._handle_attempt_error(
|
||
exc,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
max_retries=max_retries,
|
||
record_id=record_id,
|
||
attempt_count=attempt_count,
|
||
max_attempts=max_attempts,
|
||
execution_error_handler=execution_error_handler,
|
||
retry_policy=retry_policy,
|
||
candidate_record_map=candidate_record_map,
|
||
candidates=candidates,
|
||
request_id=request_id,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
)
|
||
action = outcome.action
|
||
last_status_code = outcome.last_status_code
|
||
max_retries = outcome.max_retries
|
||
if outcome.stop_result is not None:
|
||
return outcome.stop_result
|
||
if action in {FailoverAction.CONTINUE, FailoverAction.RETRY}:
|
||
consecutive_failures += 1
|
||
await self._apply_retry_pacing(
|
||
candidate=candidate,
|
||
consecutive_failures=consecutive_failures,
|
||
error=exc,
|
||
request_id=request_id,
|
||
)
|
||
|
||
# action switch: continue/ retry
|
||
if action == FailoverAction.CONTINUE:
|
||
# PRE_EXPAND: if we break early, mark remaining retries of this candidate unused.
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_candidate_remaining_retries_unused(
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_idx=candidate_index,
|
||
from_retry_idx=retry_index + 1,
|
||
retry_policy=retry_policy,
|
||
)
|
||
break
|
||
if action == FailoverAction.RETRY:
|
||
retry_index += 1
|
||
continue
|
||
|
||
# Safety: unknown action -> stop retrying this candidate.
|
||
break
|
||
|
||
# exhausted: PRE_EXPAND should not leave 'available' records behind
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_all_remaining_available_unused(candidate_record_map)
|
||
|
||
return ExecutionResult(
|
||
success=False,
|
||
error_type="AllCandidatesFailed",
|
||
error_message="All candidates exhausted",
|
||
last_status_code=last_status_code,
|
||
candidate_keys=self._get_candidate_keys(
|
||
request_id=request_id,
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
),
|
||
attempt_count=attempt_count,
|
||
)
|
||
|
||
async def _execute_pool_candidate(
|
||
self,
|
||
*,
|
||
candidate: PoolCandidate,
|
||
candidate_index: int,
|
||
attempt_func: AttemptFunc,
|
||
retry_policy: RetryPolicy,
|
||
request_id: str | None,
|
||
user_id: str | None,
|
||
api_key_id: str | None,
|
||
username: str | None,
|
||
api_key_name: str | None,
|
||
candidate_record_map: dict[tuple[int, int], str] | None,
|
||
candidate_keys_fallback: list[CandidateKey],
|
||
candidates: list[ProviderCandidate],
|
||
attempt_count: int,
|
||
consecutive_failures: int,
|
||
max_attempts: int | None,
|
||
execution_error_handler: Any,
|
||
is_cancelled: Callable[[], Awaitable[bool]] | None,
|
||
) -> tuple[ExecutionResult | None, int, int, int | None]:
|
||
"""Execute a PoolCandidate with in-pool key failover."""
|
||
last_status_code: int | None = None
|
||
retry_slots_per_key = self._get_pool_key_max_retries(candidate, retry_policy)
|
||
|
||
for key_index, pool_key in enumerate(candidate.pool_keys or []):
|
||
base_retry_index = key_index * retry_slots_per_key
|
||
cancelled_result = await self._maybe_cancel_execution(
|
||
is_cancelled=is_cancelled,
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=candidate_index,
|
||
from_retry_idx=base_retry_index,
|
||
retry_policy=retry_policy,
|
||
request_id=request_id,
|
||
attempt_count=attempt_count,
|
||
)
|
||
if cancelled_result is not None:
|
||
return cancelled_result, attempt_count, consecutive_failures, last_status_code
|
||
|
||
candidate.key = pool_key
|
||
candidate._pool_key_index = key_index
|
||
candidate.mapping_matched_model = getattr(pool_key, "_pool_mapping_matched_model", None)
|
||
|
||
if bool(getattr(pool_key, "_pool_skipped", False)):
|
||
skip_reason = str(
|
||
getattr(pool_key, "_pool_skip_reason", None)
|
||
or getattr(candidate, "skip_reason", None)
|
||
or "pool_skipped"
|
||
)
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_retry_indices_status(
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_idx=candidate_index,
|
||
retry_indices=range(
|
||
base_retry_index, base_retry_index + retry_slots_per_key
|
||
),
|
||
status="skipped",
|
||
skip_reason=skip_reason,
|
||
)
|
||
elif request_id:
|
||
await self._create_skipped_record(
|
||
request_id=request_id,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=base_retry_index,
|
||
user_id=user_id,
|
||
api_key_id=api_key_id,
|
||
username=username,
|
||
api_key_name=api_key_name,
|
||
skip_reason=skip_reason,
|
||
)
|
||
candidate_keys_fallback.append(
|
||
self._make_candidate_key(
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=base_retry_index,
|
||
status="skipped",
|
||
skip_reason=skip_reason,
|
||
)
|
||
)
|
||
continue
|
||
|
||
max_retries_for_key = retry_slots_per_key
|
||
retry_index = 0
|
||
while retry_index < max_retries_for_key:
|
||
composite_retry_index = base_retry_index + retry_index
|
||
cancelled_result = await self._maybe_cancel_execution(
|
||
is_cancelled=is_cancelled,
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=candidate_index,
|
||
from_retry_idx=composite_retry_index,
|
||
retry_policy=retry_policy,
|
||
request_id=request_id,
|
||
attempt_count=attempt_count,
|
||
)
|
||
if cancelled_result is not None:
|
||
return cancelled_result, attempt_count, consecutive_failures, last_status_code
|
||
|
||
attempt_count += 1
|
||
|
||
record_id = None
|
||
if candidate_record_map:
|
||
record_id = candidate_record_map.get((candidate_index, composite_retry_index))
|
||
if record_id is None:
|
||
# Rectify may extend retries beyond pre-created range.
|
||
record_id = candidate_record_map.get((candidate_index, base_retry_index))
|
||
if record_id is None and request_id and retry_policy.mode != RetryMode.PRE_EXPAND:
|
||
record_id = await self._ensure_record_exists(
|
||
request_id=request_id,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=composite_retry_index,
|
||
user_id=user_id,
|
||
api_key_id=api_key_id,
|
||
username=username,
|
||
api_key_name=api_key_name,
|
||
)
|
||
|
||
self._attach_attempt_context(
|
||
candidate,
|
||
candidate_index,
|
||
composite_retry_index,
|
||
record_id,
|
||
attempt_count,
|
||
max_attempts,
|
||
)
|
||
|
||
await self._mark_pending_and_commit(record_id)
|
||
|
||
try:
|
||
attempt_result = await self._execute_attempt(
|
||
candidate=candidate,
|
||
record_id=record_id,
|
||
attempt_func=attempt_func,
|
||
)
|
||
last_status_code = int(getattr(attempt_result, "http_status", 0) or 0)
|
||
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_remaining_slots_unused(
|
||
candidate_record_map=candidate_record_map,
|
||
candidates=candidates,
|
||
success_candidate_idx=candidate_index,
|
||
success_retry_idx=composite_retry_index,
|
||
retry_policy=retry_policy,
|
||
)
|
||
|
||
consecutive_failures = 0
|
||
return (
|
||
ExecutionResult(
|
||
success=True,
|
||
attempt_result=attempt_result,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=composite_retry_index,
|
||
provider_id=str(candidate.provider.id),
|
||
provider_name=str(candidate.provider.name),
|
||
endpoint_id=str(candidate.endpoint.id),
|
||
key_id=str(candidate.key.id),
|
||
candidate_keys=self._get_candidate_keys(
|
||
request_id=request_id,
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
),
|
||
attempt_count=attempt_count,
|
||
request_candidate_id=record_id,
|
||
),
|
||
attempt_count,
|
||
consecutive_failures,
|
||
last_status_code,
|
||
)
|
||
|
||
except StreamProbeError as exc:
|
||
last_status_code = exc.http_status
|
||
await self._record_attempt_failure(record_id, exc, exc.http_status)
|
||
action = FailoverAction.CONTINUE
|
||
consecutive_failures += 1
|
||
await self._apply_retry_pacing(
|
||
candidate=candidate,
|
||
consecutive_failures=consecutive_failures,
|
||
error=exc,
|
||
request_id=request_id,
|
||
)
|
||
|
||
except Exception as exc:
|
||
if _is_client_disconnected(exc):
|
||
await self._record_attempt_failure(record_id, exc, 499)
|
||
return (
|
||
await self._build_disconnected_result(
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_keys_fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
from_candidate_idx=candidate_index,
|
||
from_retry_idx=composite_retry_index + 1,
|
||
retry_policy=retry_policy,
|
||
request_id=request_id,
|
||
attempt_count=attempt_count,
|
||
),
|
||
attempt_count,
|
||
consecutive_failures,
|
||
499,
|
||
)
|
||
|
||
outcome = await self._handle_pool_attempt_error(
|
||
exc,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
key_retry_index=retry_index,
|
||
composite_retry_index=composite_retry_index,
|
||
max_retries=max_retries_for_key,
|
||
record_id=record_id,
|
||
attempt_count=attempt_count,
|
||
max_attempts=max_attempts,
|
||
execution_error_handler=execution_error_handler,
|
||
)
|
||
action = outcome.action
|
||
last_status_code = outcome.last_status_code
|
||
max_retries_for_key = min(outcome.max_retries, retry_slots_per_key)
|
||
if action in {FailoverAction.CONTINUE, FailoverAction.RETRY}:
|
||
consecutive_failures += 1
|
||
await self._apply_retry_pacing(
|
||
candidate=candidate,
|
||
consecutive_failures=consecutive_failures,
|
||
error=exc,
|
||
request_id=request_id,
|
||
)
|
||
|
||
if action == FailoverAction.CONTINUE:
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
# max_retries_for_key may have been shrunk by error handler;
|
||
# mark unused up to the *original* retry_slots_per_key to cover
|
||
# all pre-created records.
|
||
await self._mark_retry_indices_status(
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_idx=candidate_index,
|
||
retry_indices=range(
|
||
composite_retry_index + 1,
|
||
base_retry_index + retry_slots_per_key,
|
||
),
|
||
status="unused",
|
||
)
|
||
break
|
||
if action == FailoverAction.RETRY:
|
||
retry_index += 1
|
||
continue
|
||
|
||
# STOP: only stop this pool candidate; outer candidate traversal continues.
|
||
# Rationale: pool-internal STOP (from error_stop_patterns on a non-ExecutionError)
|
||
# should not terminate the entire request because other providers may still succeed.
|
||
# When handler_used=True, TaskService raises directly for true STOP semantics.
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_candidate_remaining_retries_unused(
|
||
candidate_record_map=candidate_record_map,
|
||
candidate_idx=candidate_index,
|
||
from_retry_idx=composite_retry_index + 1,
|
||
retry_policy=retry_policy,
|
||
)
|
||
return None, attempt_count, consecutive_failures, last_status_code
|
||
|
||
return None, attempt_count, consecutive_failures, last_status_code
|
||
|
||
async def _handle_pool_attempt_error(
|
||
self,
|
||
exc: Exception,
|
||
*,
|
||
candidate: ProviderCandidate,
|
||
candidate_index: int,
|
||
key_retry_index: int,
|
||
composite_retry_index: int,
|
||
max_retries: int,
|
||
record_id: str | None,
|
||
attempt_count: int,
|
||
max_attempts: int | None,
|
||
execution_error_handler: Any,
|
||
) -> AttemptErrorOutcome:
|
||
"""Handle pool attempt errors without forcing outer STOP semantics.
|
||
|
||
Args:
|
||
key_retry_index: key 内部的重试索引 (用于判断 has_retry_left)
|
||
composite_retry_index: 全局维度的重试索引 (传给 execution_error_handler,
|
||
与 candidate_record_map 对齐)
|
||
"""
|
||
return await self._classify_attempt_error(
|
||
exc,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=composite_retry_index,
|
||
has_retry_left=key_retry_index + 1 < max_retries,
|
||
max_retries=max_retries,
|
||
record_id=record_id,
|
||
attempt_count=attempt_count,
|
||
max_attempts=max_attempts,
|
||
execution_error_handler=execution_error_handler,
|
||
)
|
||
|
||
@staticmethod
|
||
def _attach_attempt_context(
|
||
candidate: ProviderCandidate,
|
||
candidate_index: int,
|
||
retry_index: int,
|
||
record_id: str | None,
|
||
attempt_count: int,
|
||
max_attempts: int | None,
|
||
) -> None:
|
||
"""Attach per-attempt context onto candidate for attempt_func (best-effort)."""
|
||
try:
|
||
setattr(candidate, "_utf_candidate_index", candidate_index)
|
||
setattr(candidate, "_utf_retry_index", retry_index)
|
||
setattr(candidate, "_utf_candidate_record_id", record_id)
|
||
setattr(candidate, "_utf_attempt_count", attempt_count)
|
||
setattr(candidate, "_utf_max_attempts", max_attempts)
|
||
except Exception:
|
||
pass
|
||
|
||
async def _execute_attempt(
|
||
self,
|
||
*,
|
||
candidate: ProviderCandidate,
|
||
record_id: str | None,
|
||
attempt_func: AttemptFunc,
|
||
) -> AttemptResult:
|
||
"""Run attempt_func with stream probe and sync failover-pattern checks.
|
||
|
||
On success records the attempt; raises StreamProbeError on failover-pattern
|
||
match or stream probe failure so the caller can handle retries uniformly.
|
||
"""
|
||
attempt_result = await attempt_func(candidate)
|
||
|
||
if attempt_result.kind == AttemptKind.STREAM:
|
||
attempt_result = await self._probe_stream_first_chunk(
|
||
attempt_result=attempt_result,
|
||
record_id=record_id,
|
||
candidate=candidate,
|
||
)
|
||
|
||
if attempt_result.kind == AttemptKind.SYNC_RESPONSE:
|
||
body = getattr(attempt_result, "response_body", None)
|
||
if body:
|
||
if isinstance(body, bytes):
|
||
body_text = body.decode("utf-8", errors="replace")
|
||
elif isinstance(body, (dict, list)):
|
||
body_text = json.dumps(body, ensure_ascii=False)
|
||
else:
|
||
body_text = str(body)
|
||
rule_action = self._check_provider_failover_rules(
|
||
candidate, is_success=True, response_text=body_text
|
||
)
|
||
if rule_action == FailoverAction.CONTINUE:
|
||
await self._record_attempt_failure(
|
||
record_id,
|
||
Exception("success_failover_pattern matched"),
|
||
200,
|
||
)
|
||
raise StreamProbeError(
|
||
"Success failover pattern matched",
|
||
http_status=200,
|
||
)
|
||
|
||
await self._record_attempt_success(record_id, attempt_result)
|
||
return attempt_result
|
||
|
||
async def _record_attempt_success(
|
||
self, record_id: str | None, attempt_result: AttemptResult
|
||
) -> None:
|
||
"""Mark attempt record as success/streaming."""
|
||
if not record_id:
|
||
return
|
||
|
||
def _do() -> None:
|
||
if attempt_result.kind == AttemptKind.STREAM:
|
||
self._update_record(
|
||
record_id,
|
||
status="streaming",
|
||
status_code=attempt_result.http_status,
|
||
)
|
||
else:
|
||
self._update_record(
|
||
record_id,
|
||
status="success",
|
||
status_code=attempt_result.http_status,
|
||
finished_at=datetime.now(timezone.utc),
|
||
)
|
||
self.db.commit()
|
||
|
||
await self._db_op(_do)
|
||
|
||
async def _record_attempt_failure(
|
||
self, record_id: str | None, exc: Exception, status_code: int | None = None
|
||
) -> None:
|
||
"""Mark attempt record as failed."""
|
||
if not record_id:
|
||
return
|
||
error_type = type(exc).__name__
|
||
error_message = self._sanitize(str(exc))
|
||
|
||
def _do() -> None:
|
||
self._update_record(
|
||
record_id,
|
||
status="failed",
|
||
status_code=status_code,
|
||
error_type=error_type,
|
||
error_message=error_message,
|
||
finished_at=datetime.now(timezone.utc),
|
||
)
|
||
self.db.commit()
|
||
|
||
await self._db_op(_do)
|
||
|
||
async def _handle_attempt_error(
|
||
self,
|
||
exc: Exception,
|
||
*,
|
||
candidate: ProviderCandidate,
|
||
candidate_index: int,
|
||
retry_index: int,
|
||
max_retries: int,
|
||
record_id: str | None,
|
||
attempt_count: int,
|
||
max_attempts: int | None,
|
||
execution_error_handler: Any,
|
||
retry_policy: RetryPolicy,
|
||
candidate_record_map: dict[tuple[int, int], str] | None,
|
||
candidates: list[ProviderCandidate],
|
||
request_id: str | None,
|
||
candidate_keys_fallback: list[CandidateKey],
|
||
) -> AttemptErrorOutcome:
|
||
"""
|
||
Handle attempt exception: delegate to external/internal handler, update records.
|
||
|
||
Returns:
|
||
AttemptErrorOutcome; stop_result is non-None only when action==STOP.
|
||
"""
|
||
outcome = await self._classify_attempt_error(
|
||
exc,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
has_retry_left=retry_index + 1 < max_retries,
|
||
max_retries=max_retries,
|
||
record_id=record_id,
|
||
attempt_count=attempt_count,
|
||
max_attempts=max_attempts,
|
||
execution_error_handler=execution_error_handler,
|
||
)
|
||
|
||
if outcome.action == FailoverAction.STOP:
|
||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||
await self._mark_remaining_slots_unused(
|
||
candidate_record_map=candidate_record_map,
|
||
candidates=candidates,
|
||
success_candidate_idx=candidate_index,
|
||
success_retry_idx=retry_index,
|
||
retry_policy=retry_policy,
|
||
)
|
||
outcome.stop_result = ExecutionResult(
|
||
success=False,
|
||
error_type=type(exc).__name__,
|
||
error_message=self._sanitize(str(exc)),
|
||
last_status_code=outcome.last_status_code or None,
|
||
candidate_keys=self._get_candidate_keys(
|
||
request_id=request_id,
|
||
fallback=candidate_keys_fallback,
|
||
candidates=candidates,
|
||
),
|
||
attempt_count=attempt_count,
|
||
)
|
||
|
||
return outcome
|
||
|
||
async def _classify_attempt_error(
|
||
self,
|
||
exc: Exception,
|
||
*,
|
||
candidate: ProviderCandidate,
|
||
candidate_index: int,
|
||
retry_index: int,
|
||
has_retry_left: bool,
|
||
max_retries: int,
|
||
record_id: str | None,
|
||
attempt_count: int,
|
||
max_attempts: int | None,
|
||
execution_error_handler: Any,
|
||
) -> AttemptErrorOutcome:
|
||
"""Classify an attempt error: delegate to external handler or internal classifier.
|
||
|
||
Returns a base AttemptErrorOutcome (without stop_result). Callers add
|
||
STOP-specific logic (e.g. PRE_EXPAND cleanup, stop_result construction) as needed.
|
||
"""
|
||
handler_used = False
|
||
action = FailoverAction.CONTINUE
|
||
if execution_error_handler is not None:
|
||
try:
|
||
from src.services.request.executor import ExecutionError as _ExecutionError
|
||
|
||
if isinstance(exc, _ExecutionError):
|
||
handler_used = True
|
||
action, new_max_retries = await execution_error_handler(
|
||
exec_err=exc,
|
||
candidate=candidate,
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
max_retries_for_candidate=max_retries,
|
||
record_id=record_id,
|
||
attempt_count=attempt_count,
|
||
max_attempts=max_attempts,
|
||
)
|
||
if new_max_retries is not None:
|
||
max_retries = max(max_retries, int(new_max_retries))
|
||
except Exception:
|
||
handler_used = False
|
||
|
||
last_status_code: int | None = None
|
||
if not handler_used:
|
||
action = await self._handle_error(
|
||
exc,
|
||
candidate=candidate,
|
||
has_retry_left=has_retry_left,
|
||
)
|
||
last_status_code = int(getattr(exc, "status_code", 0) or 0) or int(
|
||
getattr(exc, "http_status", 0) or 0
|
||
)
|
||
await self._record_attempt_failure(record_id, exc, last_status_code or None)
|
||
|
||
return AttemptErrorOutcome(
|
||
action=action,
|
||
last_status_code=last_status_code,
|
||
max_retries=max_retries,
|
||
)
|
||
|
||
def _sanitize(self, message: str, max_length: int = 200) -> str:
|
||
if not message:
|
||
return "request_failed"
|
||
return _SENSITIVE_PATTERN.sub("[REDACTED]", message)[:max_length]
|
||
|
||
def _make_candidate_key(
|
||
self,
|
||
*,
|
||
candidate: ProviderCandidate,
|
||
candidate_index: int,
|
||
retry_index: int,
|
||
status: str,
|
||
skip_reason: str | None = None,
|
||
error_type: str | None = None,
|
||
error_message: str | None = None,
|
||
status_code: int | None = None,
|
||
) -> CandidateKey:
|
||
return CandidateKey(
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
provider_id=str(candidate.provider.id),
|
||
provider_name=str(candidate.provider.name),
|
||
endpoint_id=str(candidate.endpoint.id),
|
||
key_id=str(candidate.key.id),
|
||
key_name=str(getattr(candidate.key, "name", "") or ""),
|
||
auth_type=str(getattr(candidate.key, "auth_type", "") or ""),
|
||
priority=int(getattr(candidate.key, "priority", 0) or 0),
|
||
is_cached=bool(getattr(candidate, "is_cached", False)),
|
||
status=status,
|
||
skip_reason=skip_reason,
|
||
error_type=error_type,
|
||
error_message=error_message,
|
||
status_code=status_code,
|
||
)
|
||
|
||
def _get_candidate_keys(
|
||
self,
|
||
*,
|
||
request_id: str | None,
|
||
fallback: list[CandidateKey],
|
||
candidates: list[ProviderCandidate],
|
||
) -> list[CandidateKey]:
|
||
if request_id:
|
||
try:
|
||
return self._recorder.get_candidate_keys(request_id)
|
||
except Exception as exc:
|
||
# 降级到 fallback 但记录 warning（影响审计追踪可见性）
|
||
logger.warning(
|
||
"[FailoverEngine] get_candidate_keys failed, using fallback: {}",
|
||
self._sanitize(str(exc)),
|
||
)
|
||
if fallback:
|
||
return fallback
|
||
# fallback snapshot (no DB audit)
|
||
result: list[CandidateKey] = []
|
||
for idx, cand in enumerate(candidates):
|
||
result.append(
|
||
self._make_candidate_key(
|
||
candidate=cand,
|
||
candidate_index=idx,
|
||
retry_index=0,
|
||
status="available",
|
||
)
|
||
)
|
||
return result
|
||
|
||
async def _commit_before_await(self) -> None:
|
||
def _do() -> None:
|
||
if self.db.in_transaction():
|
||
try:
|
||
self.db.commit()
|
||
except Exception:
|
||
self.db.rollback()
|
||
raise
|
||
|
||
await self._db_op(_do)
|
||
|
||
async def _mark_pending_and_commit(self, record_id: str | None) -> None:
|
||
"""Mark record as pending and commit, all within a worker thread."""
|
||
|
||
def _do() -> None:
|
||
if record_id:
|
||
self._update_record(
|
||
record_id, status="pending", started_at=datetime.now(timezone.utc)
|
||
)
|
||
if self.db.in_transaction():
|
||
try:
|
||
self.db.commit()
|
||
except Exception:
|
||
self.db.rollback()
|
||
raise
|
||
|
||
await self._db_op(_do)
|
||
|
||
def _update_record(self, record_id: str, /, **values: Any) -> None:
|
||
self.db.execute(
|
||
update(RequestCandidate).where(RequestCandidate.id == record_id).values(**values)
|
||
)
|
||
|
||
async def _ensure_record_exists(
|
||
self,
|
||
*,
|
||
request_id: str,
|
||
candidate: ProviderCandidate,
|
||
candidate_index: int,
|
||
retry_index: int,
|
||
user_id: str | None,
|
||
api_key_id: str | None,
|
||
username: str | None,
|
||
api_key_name: str | None,
|
||
) -> str:
|
||
# Create "available" record, then caller will mark pending.
|
||
extra = self._build_pool_extra_data(candidate)
|
||
provider_id = str(candidate.provider.id)
|
||
endpoint_id = str(candidate.endpoint.id)
|
||
key_id = str(candidate.key.id)
|
||
is_cached = bool(getattr(candidate, "is_cached", False))
|
||
|
||
def _do() -> str:
|
||
row = RequestCandidateService.create_candidate(
|
||
db=self.db,
|
||
request_id=request_id,
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
user_id=user_id,
|
||
api_key_id=api_key_id,
|
||
username=username,
|
||
api_key_name=api_key_name,
|
||
provider_id=provider_id,
|
||
endpoint_id=endpoint_id,
|
||
key_id=key_id,
|
||
status="available",
|
||
is_cached=is_cached,
|
||
extra_data=extra,
|
||
)
|
||
return str(row.id)
|
||
|
||
return await self._db_op(_do)
|
||
|
||
async def _create_skipped_record(
|
||
self,
|
||
*,
|
||
request_id: str,
|
||
candidate: ProviderCandidate,
|
||
candidate_index: int,
|
||
retry_index: int = 0,
|
||
user_id: str | None,
|
||
api_key_id: str | None,
|
||
username: str | None,
|
||
api_key_name: str | None,
|
||
skip_reason: str | None,
|
||
) -> str:
|
||
extra = self._build_pool_extra_data(candidate)
|
||
provider_id = str(candidate.provider.id)
|
||
endpoint_id = str(candidate.endpoint.id)
|
||
key_id = str(candidate.key.id)
|
||
is_cached = bool(getattr(candidate, "is_cached", False))
|
||
|
||
def _do() -> str:
|
||
row = RequestCandidateService.create_candidate(
|
||
db=self.db,
|
||
request_id=request_id,
|
||
candidate_index=candidate_index,
|
||
retry_index=retry_index,
|
||
user_id=user_id,
|
||
api_key_id=api_key_id,
|
||
username=username,
|
||
api_key_name=api_key_name,
|
||
provider_id=provider_id,
|
||
endpoint_id=endpoint_id,
|
||
key_id=key_id,
|
||
status="skipped",
|
||
skip_reason=skip_reason,
|
||
is_cached=is_cached,
|
||
extra_data=extra,
|
||
)
|
||
# ensure visible for subsequent recorder reads
|
||
if self.db.in_transaction():
|
||
self.db.commit()
|
||
return str(row.id)
|
||
|
||
return await self._db_op(_do)
|
||
|
||
def _build_pool_extra_data(self, candidate: ProviderCandidate) -> dict[str, Any]:
|
||
extra: dict[str, Any] = {}
|
||
pool_extra = getattr(candidate, "_pool_extra_data", None)
|
||
if isinstance(pool_extra, dict):
|
||
extra.update(pool_extra)
|
||
|
||
if isinstance(candidate, PoolCandidate):
|
||
extra["pool_group_id"] = str(getattr(candidate.provider, "id", "") or "")
|
||
extra["pool_key_index"] = int(getattr(candidate, "_pool_key_index", 0) or 0)
|
||
key_extra = getattr(candidate.key, "_pool_extra_data", None)
|
||
if isinstance(key_extra, dict):
|
||
extra.update(key_extra)
|
||
return extra
|
||
|
||
def _should_skip(
|
||
self, candidate: ProviderCandidate, skip_policy: SkipPolicy
|
||
) -> tuple[bool, str | None]:
|
||
if bool(getattr(candidate, "is_skipped", False)):
|
||
return True, str(getattr(candidate, "skip_reason", None) or "scheduler_marked")
|
||
|
||
auth_type = str(getattr(getattr(candidate, "key", None), "auth_type", "") or "api_key")
|
||
if (
|
||
skip_policy.supported_auth_types is not None
|
||
and auth_type not in skip_policy.supported_auth_types
|
||
):
|
||
return True, "unsupported_auth_type"
|
||
|
||
needs_conversion = bool(getattr(candidate, "needs_conversion", False))
|
||
if needs_conversion and not skip_policy.allow_format_conversion:
|
||
return True, "format_conversion_not_supported"
|
||
|
||
return False, None
|
||
|
||
def _get_max_retries(self, candidate: ProviderCandidate, retry_policy: RetryPolicy) -> int:
|
||
per_key_retries = self._get_pool_key_max_retries(candidate, retry_policy)
|
||
if isinstance(candidate, PoolCandidate):
|
||
key_count = len(candidate.pool_keys or []) or 1
|
||
return max(1, key_count * per_key_retries)
|
||
return per_key_retries
|
||
|
||
def _get_pool_key_max_retries(
|
||
self, candidate: ProviderCandidate, retry_policy: RetryPolicy
|
||
) -> int:
|
||
if retry_policy.mode == RetryMode.DISABLED:
|
||
return 1
|
||
if retry_policy.retry_on_cached_only and not bool(getattr(candidate, "is_cached", False)):
|
||
return 1
|
||
provider_max = getattr(getattr(candidate, "provider", None), "max_retries", None)
|
||
try:
|
||
value = int(provider_max or retry_policy.max_retries or 1)
|
||
except Exception:
|
||
value = int(retry_policy.max_retries or 1)
|
||
return max(1, value)
|
||
|
||
async def _handle_error(
|
||
self,
|
||
error: Exception,
|
||
*,
|
||
candidate: ProviderCandidate,
|
||
has_retry_left: bool,
|
||
) -> FailoverAction:
|
||
# 检查提供商级别的错误终止规则
|
||
error_text = self._extract_error_text(error)
|
||
status_code = int(getattr(error, "status_code", 0) or 0) or int(
|
||
getattr(error, "http_status", 0) or 0
|
||
)
|
||
# ExecutionError wrapping: check cause for status_code
|
||
if not status_code:
|
||
cause = getattr(error, "cause", None)
|
||
if cause is not None:
|
||
status_code = int(getattr(cause, "status_code", 0) or 0) or int(
|
||
getattr(cause, "http_status", 0) or 0
|
||
)
|
||
if error_text:
|
||
rule_action = self._check_provider_failover_rules(
|
||
candidate,
|
||
is_success=False,
|
||
response_text=error_text,
|
||
status_code=status_code or None,
|
||
)
|
||
if rule_action is not None:
|
||
return rule_action
|
||
|
||
# 默认全部转移: ErrorClassifier 结果统一映射为 CONTINUE/RETRY，不再 STOP
|
||
action = self._error_classifier.classify(error, has_retry_left=has_retry_left)
|
||
if action == ErrorAction.CONTINUE:
|
||
return FailoverAction.RETRY
|
||
return FailoverAction.CONTINUE
|
||
|
||
def _check_provider_failover_rules(
|
||
self,
|
||
candidate: ProviderCandidate,
|
||
*,
|
||
is_success: bool,
|
||
response_text: str,
|
||
status_code: int | None = None,
|
||
) -> FailoverAction | None:
|
||
"""检查提供商级别的故障转移规则。返回 None 表示无规则命中，使用默认行为。"""
|
||
config = getattr(candidate.provider, "config", None) or {}
|
||
rules = config.get("failover_rules")
|
||
if not rules or not isinstance(rules, dict):
|
||
return None
|
||
|
||
compiled = self._get_compiled_patterns(rules)
|
||
|
||
if is_success:
|
||
for regex, rule in compiled.get("success", []):
|
||
if regex.search(response_text):
|
||
logger.info(
|
||
"[FailoverEngine] 成功转移规则命中: pattern={}, provider={}",
|
||
rule.get("pattern", ""),
|
||
candidate.provider.name,
|
||
)
|
||
return FailoverAction.CONTINUE
|
||
else:
|
||
for regex, rule in compiled.get("error", []):
|
||
# 检查状态码过滤
|
||
rule_status_codes = rule.get("status_codes")
|
||
if rule_status_codes and status_code not in rule_status_codes:
|
||
continue
|
||
if regex.search(response_text):
|
||
logger.info(
|
||
"[FailoverEngine] 错误终止规则命中: pattern={}, status_code={}, provider={}",
|
||
rule.get("pattern", ""),
|
||
status_code,
|
||
candidate.provider.name,
|
||
)
|
||
return FailoverAction.STOP
|
||
|
||
return None
|
||
|
||
@staticmethod
|
||
def _get_compiled_patterns(
|
||
rules: dict[str, Any],
|
||
) -> dict[str, list[tuple[re.Pattern[str], dict[str, Any]]]]:
|
||
"""编译 failover_rules 中的正则模式。
|
||
|
||
编译结果缓存在 rules dict 的 _compiled 键上，避免每次请求都重复编译。
|
||
"""
|
||
cached = rules.get("_compiled")
|
||
if cached is not None:
|
||
return cached
|
||
|
||
result: dict[str, list[tuple[re.Pattern[str], dict[str, Any]]]] = {
|
||
"success": [],
|
||
"error": [],
|
||
}
|
||
for rule in rules.get("success_failover_patterns", []):
|
||
pattern = rule.get("pattern", "")
|
||
if pattern:
|
||
try:
|
||
result["success"].append((re.compile(pattern), rule))
|
||
except re.error:
|
||
pass
|
||
for rule in rules.get("error_stop_patterns", []):
|
||
pattern = rule.get("pattern", "")
|
||
if pattern:
|
||
try:
|
||
result["error"].append((re.compile(pattern), rule))
|
||
except re.error:
|
||
pass
|
||
rules["_compiled"] = result
|
||
return result
|
||
|
||
@staticmethod
|
||
def _extract_error_text(error: Exception) -> str:
|
||
"""从异常中提取错误响应文本。"""
|
||
# ExecutionError wrapping
|
||
cause = getattr(error, "cause", None)
|
||
if cause is not None:
|
||
error = cause
|
||
|
||
# httpx.HTTPStatusError
|
||
response = getattr(error, "response", None)
|
||
if response is not None:
|
||
try:
|
||
return response.text or ""
|
||
except Exception:
|
||
pass
|
||
|
||
# upstream_response / upstream_error attribute
|
||
for attr in ("upstream_response", "upstream_error", "error_message"):
|
||
val = getattr(error, attr, None)
|
||
if val:
|
||
return str(val)
|
||
|
||
return str(error)
|
||
|
||
async def _probe_stream_first_chunk(
|
||
self,
|
||
*,
|
||
attempt_result: AttemptResult,
|
||
record_id: str | None,
|
||
candidate: ProviderCandidate | None = None,
|
||
) -> AttemptResult:
|
||
"""
|
||
Probe first chunk for a streaming response.
|
||
|
||
Strong constraints:
|
||
- Must have timeout.
|
||
- Empty stream before first chunk is treated as probe failure (eligible for failover).
|
||
"""
|
||
assert attempt_result.kind == AttemptKind.STREAM
|
||
assert attempt_result.stream_iterator is not None
|
||
|
||
original_iterator = attempt_result.stream_iterator
|
||
try:
|
||
first_chunk = await asyncio.wait_for(
|
||
original_iterator.__anext__(),
|
||
timeout=self.STREAM_FIRST_CHUNK_TIMEOUT_SECONDS,
|
||
)
|
||
except asyncio.TimeoutError as exc:
|
||
raise StreamProbeError(
|
||
"Timeout waiting for first chunk",
|
||
http_status=attempt_result.http_status,
|
||
original_exception=exc,
|
||
) from exc
|
||
except StopAsyncIteration as exc:
|
||
raise StreamProbeError(
|
||
"Empty stream: no data received before EOF",
|
||
http_status=attempt_result.http_status,
|
||
original_exception=exc,
|
||
) from exc
|
||
except Exception as exc:
|
||
raise StreamProbeError(
|
||
f"Failed to read first chunk: {exc}",
|
||
http_status=attempt_result.http_status,
|
||
original_exception=exc,
|
||
) from exc
|
||
|
||
# Check success_failover_patterns on first chunk
|
||
if candidate is not None and first_chunk:
|
||
chunk_text = (
|
||
first_chunk.decode("utf-8", errors="replace")
|
||
if isinstance(first_chunk, bytes)
|
||
else str(first_chunk)
|
||
)
|
||
rule_action = self._check_provider_failover_rules(
|
||
candidate, is_success=True, response_text=chunk_text
|
||
)
|
||
if rule_action == FailoverAction.CONTINUE:
|
||
raise StreamProbeError(
|
||
"Success failover pattern matched in first chunk",
|
||
http_status=attempt_result.http_status,
|
||
)
|
||
|
||
wrapped = self._wrap_stream_with_finalizer(
|
||
first_chunk=first_chunk,
|
||
original_iterator=original_iterator,
|
||
record_id=record_id,
|
||
)
|
||
return AttemptResult(
|
||
kind=AttemptKind.STREAM,
|
||
http_status=attempt_result.http_status,
|
||
http_headers=attempt_result.http_headers,
|
||
stream_iterator=wrapped,
|
||
raw_response=attempt_result.raw_response,
|
||
)
|
||
|
||
def _wrap_stream_with_finalizer(
|
||
self,
|
||
*,
|
||
first_chunk: bytes,
|
||
original_iterator: AsyncIterator[bytes],
|
||
record_id: str | None,
|
||
) -> AsyncIterator[bytes]:
|
||
async def _gen() -> AsyncIterator[bytes]:
|
||
yield first_chunk
|
||
try:
|
||
async for chunk in original_iterator:
|
||
yield chunk
|
||
except Exception as exc:
|
||
# Best-effort: mark stream interrupted using a new session (stream may outlive request session).
|
||
if record_id:
|
||
self._mark_record_stream_interrupted(record_id, exc)
|
||
raise
|
||
|
||
return _gen()
|
||
|
||
def _mark_record_stream_interrupted(self, record_id: str, exc: Exception) -> None:
|
||
try:
|
||
from src.database import create_session
|
||
|
||
with create_session() as db:
|
||
db.execute(
|
||
update(RequestCandidate)
|
||
.where(RequestCandidate.id == record_id)
|
||
.values(
|
||
status="stream_interrupted",
|
||
error_type=type(exc).__name__,
|
||
error_message=self._sanitize(str(exc)),
|
||
finished_at=datetime.now(timezone.utc),
|
||
)
|
||
)
|
||
db.commit()
|
||
except Exception as inner:
|
||
logger.debug(
|
||
"[FailoverEngine] Failed to mark stream_interrupted: {}",
|
||
self._sanitize(str(inner)),
|
||
)
|
||
|
||
async def _mark_candidate_skipped(
|
||
self,
|
||
*,
|
||
candidate_record_map: dict[tuple[int, int], str],
|
||
candidate_index: int,
|
||
candidate: ProviderCandidate,
|
||
retry_policy: RetryPolicy,
|
||
skip_reason: str | None,
|
||
) -> None:
|
||
max_retries = self._get_max_retries(candidate, retry_policy)
|
||
record_ids = [
|
||
candidate_record_map[candidate_index, ri]
|
||
for ri in range(max_retries)
|
||
if (candidate_index, ri) in candidate_record_map
|
||
]
|
||
if not record_ids:
|
||
return
|
||
|
||
def _do() -> None:
|
||
now = datetime.now(timezone.utc)
|
||
for rid in record_ids:
|
||
self._update_record(rid, status="skipped", skip_reason=skip_reason, finished_at=now)
|
||
self.db.commit()
|
||
|
||
await self._db_op(_do)
|
||
|
||
async def _mark_remaining_slots_unused(
|
||
self,
|
||
*,
|
||
candidate_record_map: dict[tuple[int, int], str],
|
||
candidates: list[ProviderCandidate],
|
||
success_candidate_idx: int,
|
||
success_retry_idx: int,
|
||
retry_policy: RetryPolicy,
|
||
) -> None:
|
||
record_ids: list[str] = []
|
||
for candidate_idx, cand in enumerate(candidates):
|
||
max_retries = self._get_max_retries(cand, retry_policy)
|
||
for retry_idx in range(max_retries):
|
||
if candidate_idx < success_candidate_idx:
|
||
continue
|
||
if candidate_idx == success_candidate_idx and retry_idx <= success_retry_idx:
|
||
continue
|
||
record_id = candidate_record_map.get((candidate_idx, retry_idx))
|
||
if record_id:
|
||
record_ids.append(record_id)
|
||
|
||
if not record_ids:
|
||
return
|
||
|
||
def _do() -> None:
|
||
now = datetime.now(timezone.utc)
|
||
for rid in record_ids:
|
||
self._update_record(rid, status="unused", finished_at=now)
|
||
self.db.commit()
|
||
|
||
await self._db_op(_do)
|
||
|
||
async def _mark_candidate_remaining_retries_unused(
|
||
self,
|
||
*,
|
||
candidate_record_map: dict[tuple[int, int], str],
|
||
candidate_idx: int,
|
||
from_retry_idx: int,
|
||
retry_policy: RetryPolicy,
|
||
) -> None:
|
||
# Only meaningful for PRE_EXPAND.
|
||
upper = max(
|
||
(ri for (ci, ri) in candidate_record_map.keys() if ci == candidate_idx),
|
||
default=retry_policy.max_retries - 1,
|
||
)
|
||
record_ids = [
|
||
candidate_record_map[candidate_idx, ri]
|
||
for ri in range(from_retry_idx, upper + 1)
|
||
if (candidate_idx, ri) in candidate_record_map
|
||
]
|
||
if not record_ids:
|
||
return
|
||
|
||
def _do() -> None:
|
||
now = datetime.now(timezone.utc)
|
||
for rid in record_ids:
|
||
self._update_record(rid, status="unused", finished_at=now)
|
||
self.db.commit()
|
||
|
||
await self._db_op(_do)
|
||
|
||
async def _mark_retry_indices_status(
|
||
self,
|
||
*,
|
||
candidate_record_map: dict[tuple[int, int], str],
|
||
candidate_idx: int,
|
||
retry_indices: range,
|
||
status: str,
|
||
skip_reason: str | None = None,
|
||
) -> None:
|
||
record_ids = [
|
||
candidate_record_map[candidate_idx, ri]
|
||
for ri in retry_indices
|
||
if (candidate_idx, ri) in candidate_record_map
|
||
]
|
||
if not record_ids:
|
||
return
|
||
|
||
def _do() -> None:
|
||
now = datetime.now(timezone.utc)
|
||
for rid in record_ids:
|
||
values: dict[str, Any] = {"status": status, "finished_at": now}
|
||
if status == "skipped":
|
||
values["skip_reason"] = skip_reason
|
||
self._update_record(rid, **values)
|
||
self.db.commit()
|
||
|
||
await self._db_op(_do)
|
||
|
||
async def _mark_all_remaining_available_unused(
|
||
self, candidate_record_map: dict[tuple[int, int], str]
|
||
) -> None:
|
||
# As a safety net: do not leave available records behind in PRE_EXPAND mode.
|
||
ids = list(candidate_record_map.values())
|
||
if not ids:
|
||
return
|
||
|
||
def _do() -> None:
|
||
try:
|
||
now = datetime.now(timezone.utc)
|
||
self.db.execute(
|
||
update(RequestCandidate)
|
||
.where(RequestCandidate.id.in_(ids))
|
||
.where(RequestCandidate.status == "available")
|
||
.values(status="unused", finished_at=now)
|
||
)
|
||
self.db.commit()
|
||
except Exception:
|
||
self.db.rollback()
|
||
raise
|
||
|
||
await self._db_op(_do)
|