i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-09 04:30:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(stability): 健康监控 DB 操作异步化，防止 worker 超时崩溃

Browse Source 
- executor: record_success 通过 asyncio.to_thread offload 到线程池
- error_handler: 3 处 record_failure 同样 offload 到线程池
- sync_execute: 设置 expire_on_commit=False 防止 commit 后 ORM 懒加载
- handler_adapter_base: 归一化 check_endpoint 的 base_url 输入
This commit is contained in:
fawney19
2026-03-12 16:17:22 +08:00
parent 127b4e11de
commit 8d69f72e2a
6 changed files with 519 additions and 288 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/handlers/base/handler_adapter_base.py
+33 -4
View File
@@ -321,6 +321,22 @@ class HandlerAdapterBase(ApiAdapter):
_ = base_url, provider_type
return build_test_request_body(cls.FORMAT_ID, request_data)
@staticmethod
def _normalize_test_base_url(base_url: Any) -> str:
"""归一化 test-model 场景传入的 base_url。"""
if isinstance(base_url, str):
normalized = base_url.strip()
if normalized:
return normalized
elif isinstance(base_url, dict):
for key in ("base_url", "url"):
value = base_url.get(key)
if isinstance(value, str) and value.strip():
logger.debug("[check_endpoint] 兼容字典形式的 base_url 输入: key={}", key)
return value.strip()
raise TypeError("base_url must be a non-empty string or a dict containing 'base_url'/'url'")
@classmethod
async def check_endpoint(
cls,
@@ -359,6 +375,7 @@ class HandlerAdapterBase(ApiAdapter):
from src.core.api_format.headers import HeaderBuilder
from src.core.provider_types import ProviderType
normalized_base_url = cls._normalize_test_base_url(base_url)
is_antigravity = provider_type == ProviderType.ANTIGRAVITY
is_gemini_cli = provider_type == ProviderType.GEMINI_CLI
is_vertex = provider_type == ProviderType.VERTEX_AI
@@ -376,7 +393,9 @@ class HandlerAdapterBase(ApiAdapter):
_kiro_cfg = KiroAuthConfig.from_dict(decrypted_auth_config or {})
region = _kiro_cfg.effective_api_region()
effective_base_url = (
base_url.replace("{region}", region) if "{region}" in base_url else base_url
normalized_base_url.replace("{region}", region)
if "{region}" in normalized_base_url
else normalized_base_url
)
url = f"{str(effective_base_url).rstrip('/')}{KIRO_GENERATE_ASSISTANT_PATH}"
elif is_antigravity:
@@ -422,11 +441,17 @@ class HandlerAdapterBase(ApiAdapter):
)
else:
url = cls.build_endpoint_url(
base_url, request_data, model_name, provider_type=provider_type
normalized_base_url,
request_data,
model_name,
provider_type=provider_type,
)
# ---- Headers ----
cli_extra = cls.get_cli_extra_headers(base_url=base_url, provider_type=provider_type)
cli_extra = cls.get_cli_extra_headers(
base_url=normalized_base_url,
provider_type=provider_type,
)
merged_extra = dict(extra_headers) if extra_headers else {}
merged_extra.update(cli_extra)
@@ -480,7 +505,11 @@ class HandlerAdapterBase(ApiAdapter):
headers[auth_header_name] = f"Bearer {api_key}"
# ---- Body ----
body = cls.build_request_body(request_data, base_url=base_url, provider_type=provider_type)
body = cls.build_request_body(
request_data,
base_url=normalized_base_url,
provider_type=provider_type,
)
if body_rules:
body = apply_body_rules(body, body_rules)
src/services/orchestration/error_handler.py
+6 -3
View File
@@ -156,7 +156,8 @@ class ErrorHandlerService:
affinity_key, client_format_str, global_model_id, endpoint, key
)
if key:
health_monitor.record_failure(
await asyncio.to_thread(
health_monitor.record_failure,
db=self.db,
key_id=str(key.id),
api_format=provider_format_str,
@@ -224,7 +225,8 @@ class ErrorHandlerService:
# 记录健康失败
if key:
health_monitor.record_failure(
await asyncio.to_thread(
health_monitor.record_failure,
db=self.db,
key_id=str(key.id),
api_format=provider_format_str,
@@ -268,7 +270,8 @@ class ErrorHandlerService:
# 记录健康失败
if key:
health_monitor.record_failure(
await asyncio.to_thread(
health_monitor.record_failure,
db=self.db,
key_id=str(key.id),
api_format=provider_format_str,
src/services/request/executor.py
+3 -1
View File
@@ -4,6 +4,7 @@
from __future__ import annotations
import asyncio
import math
import time
from collections.abc import Callable
@@ -174,7 +175,8 @@ class RequestExecutor:
client_format_str = normalize_endpoint_signature(api_format)
health_format = provider_format_str or client_format_str
health_monitor.record_success(
await asyncio.to_thread(
health_monitor.record_success,
db=self.db,
key_id=key.id,
api_format=health_format,
src/services/task/execute/sync_execute.py
+297 -280
View File
@@ -84,313 +84,330 @@ class SyncTaskExecutionService:
if not request_id:
request_id = str(uuid4())
# Build execution components (mirrors pre-Phase-3 initialization)
priority_mode = SystemConfigService.get_config(
self.db,
"provider_priority_mode",
CacheAwareScheduler.PRIORITY_MODE_PROVIDER,
)
scheduling_mode = SystemConfigService.get_config(
self.db,
"scheduling_mode",
CacheAwareScheduler.SCHEDULING_MODE_CACHE_AFFINITY,
)
cache_scheduler = await get_cache_aware_scheduler(
self.redis,
priority_mode=priority_mode,
scheduling_mode=scheduling_mode,
)
# Ensure cache_scheduler inner state is ready
await cache_scheduler._ensure_initialized()
# IMPORTANT:
# This SYNC path awaits upstream HTTP work while the failover engine may commit
# candidate audit rows between attempts. SQLAlchemy's default
# expire_on_commit=True would expire provider/endpoint/key ORM objects and can
# trigger an unexpected lazy DB reload later in error handling (for example when
# reading candidate.provider.config for failover_rules after a timeout).
#
# Keep already-loaded candidate objects resident in memory for the duration of
# the request, mirroring the async submit path.
original_expire_on_commit = getattr(self.db, "expire_on_commit", True)
self.db.expire_on_commit = False
concurrency_manager = await get_concurrency_manager()
adaptive_manager = get_adaptive_rpm_manager()
request_executor = RequestExecutor(
db=self.db,
concurrency_manager=concurrency_manager,
adaptive_manager=adaptive_manager,
)
candidate_resolver = CandidateResolver(
db=self.db,
cache_scheduler=cache_scheduler,
)
error_classifier = ErrorClassifier(
db=self.db,
cache_scheduler=cache_scheduler,
adaptive_manager=adaptive_manager,
)
request_dispatcher = RequestDispatcher(
db=self.db,
request_executor=request_executor,
cache_scheduler=cache_scheduler,
)
affinity_key = str(user_api_key.id)
user_id = str(user_api_key.user_id)
api_format_norm = normalize_endpoint_signature(api_format)
username_snapshot = None
api_key_name_snapshot = getattr(user_api_key, "name", None)
# Keep pending usage creation behavior consistent with previous behavior
try:
user = self.db.query(User).filter(User.id == user_api_key.user_id).first()
username_snapshot = getattr(user, "username", None) if user else None
UsageService.create_pending_usage(
db=self.db,
request_id=request_id,
user=user,
api_key=user_api_key,
model=model_name,
is_stream=is_stream,
api_format=api_format_norm,
request_headers=request_headers,
request_body=request_body,
# Build execution components (mirrors pre-Phase-3 initialization)
priority_mode = SystemConfigService.get_config(
self.db,
"provider_priority_mode",
CacheAwareScheduler.PRIORITY_MODE_PROVIDER,
)
except Exception as exc:
logger.warning("创建 pending 使用记录失败: {}", str(exc))
scheduling_mode = SystemConfigService.get_config(
self.db,
"scheduling_mode",
CacheAwareScheduler.SCHEDULING_MODE_CACHE_AFFINITY,
)
cache_scheduler = await get_cache_aware_scheduler(
self.redis,
priority_mode=priority_mode,
scheduling_mode=scheduling_mode,
)
# Ensure cache_scheduler inner state is ready
await cache_scheduler._ensure_initialized()
all_candidates, global_model_id = await candidate_resolver.fetch_candidates(
api_format=api_format_norm,
model_name=model_name,
affinity_key=affinity_key,
user_api_key=user_api_key,
request_id=request_id,
is_stream=is_stream,
capability_requirements=capability_requirements,
preferred_key_ids=preferred_key_ids,
request_body=request_body,
)
concurrency_manager = await get_concurrency_manager()
adaptive_manager = get_adaptive_rpm_manager()
request_executor = RequestExecutor(
db=self.db,
concurrency_manager=concurrency_manager,
adaptive_manager=adaptive_manager,
)
candidate_resolver = CandidateResolver(
db=self.db,
cache_scheduler=cache_scheduler,
)
error_classifier = ErrorClassifier(
db=self.db,
cache_scheduler=cache_scheduler,
adaptive_manager=adaptive_manager,
)
request_dispatcher = RequestDispatcher(
db=self.db,
request_executor=request_executor,
cache_scheduler=cache_scheduler,
)
# 号池排序涉及大量 Redis 操作，提前释放 DB 连接避免连接池压力
from src.services.scheduling.utils import release_db_connection_before_await
affinity_key = str(user_api_key.id)
user_id = str(user_api_key.user_id)
api_format_norm = normalize_endpoint_signature(api_format)
username_snapshot = None
api_key_name_snapshot = getattr(user_api_key, "name", None)
release_db_connection_before_await(self.db)
# Account Pool: reorder candidates for claude_code providers.
all_candidates, pool_traces = await self._pool_ops.apply_pool_reorder(
all_candidates, request_body=request_body
)
candidate_record_map = await candidate_resolver.create_candidate_records_async(
all_candidates=all_candidates,
request_id=request_id,
user_id=user_id,
user_api_key=user_api_key,
required_capabilities=capability_requirements,
)
max_attempts = candidate_resolver.count_total_attempts(all_candidates)
# Keep behavior consistent with previous behavior: last_candidate is updated even if skipped.
execution_state = SyncExecutionState(
candidate_record_map=candidate_record_map,
request_body_ref=request_body_ref,
last_candidate=all_candidates[-1] if all_candidates else None,
)
async def _attempt(candidate: Any) -> AttemptResult:
execution_state.touch_candidate(candidate)
candidate_index = int(getattr(candidate, "_utf_candidate_index", -1))
retry_index = int(getattr(candidate, "_utf_retry_index", 0))
candidate_record_id = str(getattr(candidate, "_utf_candidate_record_id", "") or "")
attempt_counter = int(getattr(candidate, "_utf_attempt_count", 0))
max_attempts_local = int(getattr(candidate, "_utf_max_attempts", max_attempts))
# Safety net: if record_id missing, create an "available" record on-demand.
if not candidate_record_id:
from src.services.scheduling.schemas import PoolCandidate
pool_extra = (
getattr(candidate.key, "_pool_extra_data", None)
if isinstance(getattr(candidate.key, "_pool_extra_data", None), dict)
else {}
)
extra_data: dict[str, Any] = {
"needs_conversion": bool(getattr(candidate, "needs_conversion", False)),
"provider_api_format": getattr(candidate, "provider_api_format", None) or None,
"mapping_matched_model": getattr(candidate, "mapping_matched_model", None)
or None,
**pool_extra,
}
if isinstance(candidate, PoolCandidate):
extra_data["pool_group_id"] = str(candidate.provider.id)
extra_data["pool_key_index"] = int(
getattr(candidate, "_pool_key_index", 0) or 0
)
created = RequestCandidateService.create_candidate(
# Keep pending usage creation behavior consistent with previous behavior
try:
user = self.db.query(User).filter(User.id == user_api_key.user_id).first()
username_snapshot = getattr(user, "username", None) if user else None
UsageService.create_pending_usage(
db=self.db,
request_id=request_id,
candidate_index=candidate_index,
retry_index=retry_index,
user_id=user_id,
api_key_id=str(user_api_key.id),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
provider_id=str(candidate.provider.id),
endpoint_id=str(candidate.endpoint.id),
key_id=str(candidate.key.id),
status="available",
is_cached=bool(getattr(candidate, "is_cached", False)),
extra_data=extra_data,
)
candidate_record_id = str(created.id)
execution_state.candidate_record_map[(candidate_index, retry_index)] = (
candidate_record_id
user=user,
api_key=user_api_key,
model=model_name,
is_stream=is_stream,
api_format=api_format_norm,
request_headers=request_headers,
request_body=request_body,
)
except Exception as exc:
logger.warning("创建 pending 使用记录失败: {}", str(exc))
(
response,
_provider_name,
attempt_id,
_provider_id,
_endpoint_id,
_key_id,
_first_byte_time_ms,
) = await request_dispatcher.dispatch(
candidate=candidate,
candidate_index=candidate_index,
retry_index=retry_index,
candidate_record_id=candidate_record_id,
user_api_key=user_api_key,
user_id=user_id,
request_func=request_func,
request_id=request_id,
all_candidates, global_model_id = await candidate_resolver.fetch_candidates(
api_format=api_format_norm,
model_name=model_name,
affinity_key=affinity_key,
global_model_id=global_model_id,
attempt_counter=attempt_counter,
max_attempts=max_attempts_local,
user_api_key=user_api_key,
request_id=request_id,
is_stream=is_stream,
)
_ = (
attempt_id,
_provider_name,
_provider_id,
_endpoint_id,
_key_id,
_first_byte_time_ms,
capability_requirements=capability_requirements,
preferred_key_ids=preferred_key_ids,
request_body=request_body,
)
# Account Pool: on success, update sticky binding + LRU.
await self._pool_ops.pool_on_success(candidate, request_body)
# 号池排序涉及大量 Redis 操作，提前释放 DB 连接避免连接池压力
from src.services.scheduling.utils import release_db_connection_before_await
if is_stream:
release_db_connection_before_await(self.db)
# Account Pool: reorder candidates for claude_code providers.
all_candidates, pool_traces = await self._pool_ops.apply_pool_reorder(
all_candidates, request_body=request_body
)
candidate_record_map = await candidate_resolver.create_candidate_records_async(
all_candidates=all_candidates,
request_id=request_id,
user_id=user_id,
user_api_key=user_api_key,
required_capabilities=capability_requirements,
)
max_attempts = candidate_resolver.count_total_attempts(all_candidates)
# Keep behavior consistent with previous behavior: last_candidate is updated even if skipped.
execution_state = SyncExecutionState(
candidate_record_map=candidate_record_map,
request_body_ref=request_body_ref,
last_candidate=all_candidates[-1] if all_candidates else None,
)
async def _attempt(candidate: Any) -> AttemptResult:
execution_state.touch_candidate(candidate)
candidate_index = int(getattr(candidate, "_utf_candidate_index", -1))
retry_index = int(getattr(candidate, "_utf_retry_index", 0))
candidate_record_id = str(getattr(candidate, "_utf_candidate_record_id", "") or "")
attempt_counter = int(getattr(candidate, "_utf_attempt_count", 0))
max_attempts_local = int(getattr(candidate, "_utf_max_attempts", max_attempts))
# Safety net: if record_id missing, create an "available" record on-demand.
if not candidate_record_id:
from src.services.scheduling.schemas import PoolCandidate
pool_extra = (
getattr(candidate.key, "_pool_extra_data", None)
if isinstance(getattr(candidate.key, "_pool_extra_data", None), dict)
else {}
)
extra_data: dict[str, Any] = {
"needs_conversion": bool(getattr(candidate, "needs_conversion", False)),
"provider_api_format": getattr(candidate, "provider_api_format", None)
or None,
"mapping_matched_model": getattr(candidate, "mapping_matched_model", None)
or None,
**pool_extra,
}
if isinstance(candidate, PoolCandidate):
extra_data["pool_group_id"] = str(candidate.provider.id)
extra_data["pool_key_index"] = int(
getattr(candidate, "_pool_key_index", 0) or 0
)
created = 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=str(user_api_key.id),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
provider_id=str(candidate.provider.id),
endpoint_id=str(candidate.endpoint.id),
key_id=str(candidate.key.id),
status="available",
is_cached=bool(getattr(candidate, "is_cached", False)),
extra_data=extra_data,
)
candidate_record_id = str(created.id)
execution_state.candidate_record_map[(candidate_index, retry_index)] = (
candidate_record_id
)
(
response,
_provider_name,
attempt_id,
_provider_id,
_endpoint_id,
_key_id,
_first_byte_time_ms,
) = await request_dispatcher.dispatch(
candidate=candidate,
candidate_index=candidate_index,
retry_index=retry_index,
candidate_record_id=candidate_record_id,
user_api_key=user_api_key,
user_id=user_id,
request_func=request_func,
request_id=request_id,
api_format=api_format_norm,
model_name=model_name,
affinity_key=affinity_key,
global_model_id=global_model_id,
attempt_counter=attempt_counter,
max_attempts=max_attempts_local,
is_stream=is_stream,
)
_ = (
attempt_id,
_provider_name,
_provider_id,
_endpoint_id,
_key_id,
_first_byte_time_ms,
)
# Account Pool: on success, update sticky binding + LRU.
await self._pool_ops.pool_on_success(candidate, request_body)
if is_stream:
return AttemptResult(
kind=AttemptKind.STREAM,
http_status=200,
http_headers={},
stream_iterator=response,
)
return AttemptResult(
kind=AttemptKind.STREAM,
kind=AttemptKind.SYNC_RESPONSE,
http_status=200,
http_headers={},
stream_iterator=response,
response_body=response,
)
return AttemptResult(
kind=AttemptKind.SYNC_RESPONSE,
http_status=200,
http_headers={},
response_body=response,
)
async def _handle_exec_err(
*,
exec_err: Any,
candidate: Any,
candidate_index: int,
retry_index: int,
max_retries_for_candidate: int,
record_id: str | None,
attempt_count: int,
max_attempts: int | None,
) -> tuple[FailoverAction, int | None]:
execution_state.track_execution_error(exec_err=exec_err, candidate=candidate)
# Fall back to retry 0 record if needed (rectify may extend retries).
candidate_record_id = execution_state.resolve_candidate_record_id(
candidate_index=candidate_index,
record_id=record_id,
)
async def _handle_exec_err(
*,
exec_err: Any,
candidate: Any,
candidate_index: int,
retry_index: int,
max_retries_for_candidate: int,
record_id: str | None,
attempt_count: int,
max_attempts: int | None,
) -> tuple[FailoverAction, int | None]:
execution_state.track_execution_error(exec_err=exec_err, candidate=candidate)
# Fall back to retry 0 record if needed (rectify may extend retries).
candidate_record_id = execution_state.resolve_candidate_record_id(
candidate_index=candidate_index,
record_id=record_id,
)
raw_action = await self._error_ops.handle_candidate_error(
exec_err=exec_err,
candidate=candidate,
candidate_record_id=candidate_record_id,
retry_index=retry_index,
max_retries_for_candidate=max_retries_for_candidate,
affinity_key=affinity_key,
api_format=api_format_norm,
global_model_id=global_model_id,
request_id=request_id,
attempt=attempt_count,
max_attempts=int(max_attempts or 0),
request_body_ref=request_body_ref,
error_classifier=error_classifier,
)
action = classify_candidate_error_action(raw_action)
if action == CandidateErrorAction.RAISE_ERROR:
execution_state.raise_classified_error(
fallback_error=exec_err,
failure_ops=self._failure_ops,
model_name=model_name,
raw_action = await self._error_ops.handle_candidate_error(
exec_err=exec_err,
candidate=candidate,
candidate_record_id=candidate_record_id,
retry_index=retry_index,
max_retries_for_candidate=max_retries_for_candidate,
affinity_key=affinity_key,
api_format=api_format_norm,
global_model_id=global_model_id,
request_id=request_id,
attempt=attempt_count,
max_attempts=int(max_attempts or 0),
request_body_ref=request_body_ref,
error_classifier=error_classifier,
)
action = classify_candidate_error_action(raw_action)
return resolve_execution_error_transition(
action=action,
state=execution_state,
max_retries_for_candidate=max_retries_for_candidate,
retry_index=retry_index,
).as_failover_tuple()
if action == CandidateErrorAction.RAISE_ERROR:
execution_state.raise_classified_error(
fallback_error=exec_err,
failure_ops=self._failure_ops,
model_name=model_name,
api_format=api_format_norm,
)
engine = FailoverEngine(
self.db,
error_classifier=error_classifier,
recorder=self._recorder,
)
result = await engine.execute(
candidates=all_candidates,
attempt_func=_attempt,
retry_policy=RetryPolicy.for_sync_task(),
skip_policy=SkipPolicy(),
request_id=request_id,
user_id=user_id,
api_key_id=str(user_api_key.id),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
candidate_record_map=candidate_record_map,
max_attempts=max_attempts,
execution_error_handler=_handle_exec_err,
)
return resolve_execution_error_transition(
action=action,
state=execution_state,
max_retries_for_candidate=max_retries_for_candidate,
retry_index=retry_index,
).as_failover_tuple()
if result.success:
# Build pool scheduling summary from traces collected during reorder.
if pool_traces and result.key_id:
try:
attempted_key_ids: set[str] = set()
for ck in result.candidate_keys or []:
status = str(getattr(ck, "status", "") or "").strip().lower()
if status in {"", "available", "pending", "skipped", "unused"}:
continue
kid = getattr(ck, "key_id", None)
if isinstance(kid, str) and kid:
attempted_key_ids.add(kid)
if not attempted_key_ids:
attempted_key_ids.add(str(result.key_id))
engine = FailoverEngine(
self.db,
error_classifier=error_classifier,
recorder=self._recorder,
)
result = await engine.execute(
candidates=all_candidates,
attempt_func=_attempt,
retry_policy=RetryPolicy.for_sync_task(),
skip_policy=SkipPolicy(),
request_id=request_id,
user_id=user_id,
api_key_id=str(user_api_key.id),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
candidate_record_map=candidate_record_map,
max_attempts=max_attempts,
execution_error_handler=_handle_exec_err,
)
for pt in pool_traces:
summary = pt.build_summary(
result.key_id,
attempted_key_ids=attempted_key_ids,
)
if summary:
result.pool_summary = summary
break
except Exception:
pass
return result
if result.success:
# Build pool scheduling summary from traces collected during reorder.
if pool_traces and result.key_id:
try:
attempted_key_ids: set[str] = set()
for ck in result.candidate_keys or []:
status = str(getattr(ck, "status", "") or "").strip().lower()
if status in {"", "available", "pending", "skipped", "unused"}:
continue
kid = getattr(ck, "key_id", None)
if isinstance(kid, str) and kid:
attempted_key_ids.add(kid)
if not attempted_key_ids:
attempted_key_ids.add(str(result.key_id))
self._failure_ops.raise_all_failed_exception(
request_id,
max_attempts,
execution_state.last_candidate,
model_name,
api_format_norm,
execution_state.last_error,
)
for pt in pool_traces:
summary = pt.build_summary(
result.key_id,
attempted_key_ids=attempted_key_ids,
)
if summary:
result.pool_summary = summary
break
except Exception:
pass
return result
self._failure_ops.raise_all_failed_exception(
request_id,
max_attempts,
execution_state.last_candidate,
model_name,
api_format_norm,
execution_state.last_error,
)
finally:
self.db.expire_on_commit = original_expire_on_commit
tests/services/test_sync_execute_expire_on_commit.py
+118
View File
@@ -0,0 +1,118 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.services.task.core.schema import ExecutionResult
from src.services.task.execute.sync_execute import SyncTaskExecutionService
@pytest.mark.asyncio
async def test_execute_sync_unified_temporarily_disables_expire_on_commit(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
db.expire_on_commit = True
db.query.return_value.filter.return_value.first.return_value = SimpleNamespace(username="alice")
pool_ops = MagicMock()
pool_ops.apply_pool_reorder = AsyncMock(return_value=([], []))
service = SyncTaskExecutionService(
db,
None,
recorder=MagicMock(),
pool_ops=pool_ops,
error_ops=MagicMock(),
failure_ops=MagicMock(),
)
cache_scheduler = SimpleNamespace(_ensure_initialized=AsyncMock())
class _StubCandidateResolver:
def __init__(self, db: object, cache_scheduler: object) -> None:
self.db = db
self.cache_scheduler = cache_scheduler
async def fetch_candidates(self, **kwargs: object) -> tuple[list[object], str]:
return [], "gpt-4.1"
async def create_candidate_records_async(
self, **kwargs: object
) -> dict[tuple[int, int], str]:
return {}
def count_total_attempts(self, _all_candidates: list[object]) -> int:
return 0
class _StubFailoverEngine:
def __init__(self, db: object, **kwargs: object) -> None:
self.db = db
async def execute(self, **kwargs: object) -> ExecutionResult:
assert getattr(self.db, "expire_on_commit") is False
return ExecutionResult(success=True)
monkeypatch.setattr(
"src.services.task.execute.sync_execute.SystemConfigService.get_config",
lambda _db, _key, default=None: default,
)
monkeypatch.setattr(
"src.services.task.execute.sync_execute.get_cache_aware_scheduler",
AsyncMock(return_value=cache_scheduler),
)
monkeypatch.setattr(
"src.services.task.execute.sync_execute.CandidateResolver",
_StubCandidateResolver,
)
monkeypatch.setattr(
"src.services.task.execute.sync_execute.ErrorClassifier",
lambda **kwargs: MagicMock(),
)
monkeypatch.setattr(
"src.services.task.execute.sync_execute.RequestDispatcher",
lambda **kwargs: MagicMock(),
)
monkeypatch.setattr(
"src.services.task.execute.sync_execute.FailoverEngine",
_StubFailoverEngine,
)
monkeypatch.setattr(
"src.services.task.execute.sync_execute.UsageService.create_pending_usage",
lambda **kwargs: None,
)
monkeypatch.setattr(
"src.services.scheduling.utils.release_db_connection_before_await",
lambda _db: None,
)
monkeypatch.setattr(
"src.services.rate_limit.concurrency_manager.get_concurrency_manager",
AsyncMock(return_value=MagicMock()),
)
monkeypatch.setattr(
"src.services.rate_limit.adaptive_rpm.get_adaptive_rpm_manager",
lambda: MagicMock(),
)
monkeypatch.setattr(
"src.services.request.executor.RequestExecutor",
lambda **kwargs: MagicMock(),
)
result = await service.execute_sync_unified(
api_format="openai_chat",
model_name="gpt-4.1",
user_api_key=SimpleNamespace(id="key-1", user_id="user-1", name="default-key"),
request_func=AsyncMock(),
request_id="req-1",
is_stream=True,
capability_requirements=None,
preferred_key_ids=None,
request_body_ref=None,
request_headers=None,
request_body=None,
)
assert result.success is True
assert db.expire_on_commit is True
tests/unit/test_handler_adapter_base.py
+62
View File
@@ -0,0 +1,62 @@
import pytest
from src.api.handlers.claude.adapter import ClaudeChatAdapter
@pytest.mark.asyncio
async def test_check_endpoint_accepts_base_url_dict(monkeypatch: pytest.MonkeyPatch) -> None:
captured: dict[str, object] = {}
async def fake_run_endpoint_check(**kwargs):
captured.update(kwargs)
return {"status_code": 200, "headers": {}, "response_time_ms": 1, "request_id": "test"}
monkeypatch.setattr(
"src.api.handlers.base.endpoint_checker.run_endpoint_check",
fake_run_endpoint_check,
)
await ClaudeChatAdapter.check_endpoint(
client=None,
base_url={"base_url": "https://api.anthropic.com"},
api_key="test-key",
request_data={
"model": "claude-sonnet-4-5-20250929",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 32,
"stream": False,
},
)
assert captured["url"] == "https://api.anthropic.com/v1/messages"
assert isinstance(captured["json_body"], dict)
@pytest.mark.asyncio
async def test_check_endpoint_accepts_url_key_in_base_url_dict(
monkeypatch: pytest.MonkeyPatch,
) -> None:
captured: dict[str, object] = {}
async def fake_run_endpoint_check(**kwargs):
captured.update(kwargs)
return {"status_code": 200, "headers": {}, "response_time_ms": 1, "request_id": "test"}
monkeypatch.setattr(
"src.api.handlers.base.endpoint_checker.run_endpoint_check",
fake_run_endpoint_check,
)
await ClaudeChatAdapter.check_endpoint(
client=None,
base_url={"url": "https://api.anthropic.com/v1"},
api_key="test-key",
request_data={
"model": "claude-sonnet-4-5-20250929",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 32,
"stream": False,
},
)
assert captured["url"] == "https://api.anthropic.com/v1/messages"