mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat: 流式空闲超时、健康监控查询优化、限流桶内存上限与维护清理修复
Close #233 Co-authored-by: AAEE86 <ppk0227@hotmail.com> - cli_monitor_mixin: 引入 STREAM_IDLE_TIMEOUT_SECONDS(可通过环境变量配置), 流传输开始后若超出空闲窗口无新 chunk 则提前取消并返回 504,避免长时间挂起 - stream_context: 新增 managed_recorded_bodies 上下文管理器,确保 chunks 在 telemetry 完成后及时释放;stream_telemetry 使用该接口统一管理 response body 构建 - health endpoint: 将状态聚合改为 GROUP BY 直接统计,事件列表按 api_format 单独查询,避免单次 limit 拉取大量记录导致的遗漏与性能问题;同时过滤不活跃 provider/endpoint,与公开健康接口保持一致 - endpoint health service: 修正时间线数据按 endpoint_id 而非 key_id 聚合 - token_bucket: 引入 max_buckets/bucket_expiry 上限与定时清理,防止内存无限增长; 修复 refill_rate=0 时 get_reset_time 除零异常;新增 _is_unlimited_rate_limit 判断 - maintenance_scheduler: 调整清理顺序(先删整行再按窗口清理),新增 newer_than 边界参数,避免同一行在同一轮中被重复改写 - sync_execute: 新增 create_pending_usage 开关,允许已预创建记录的调用方跳过重复创建 - quota_reader / provider_ops balance: 小幅修复与健壮性提升 - Dockerfile: 添加 MALLOC_ARENA_MAX=2 环境变量以降低 gunicorn worker RSS - 补充相关测试覆盖
This commit is contained in:
@@ -36,6 +36,22 @@ def test_codex_reader_preserves_summary_formats() -> None:
|
||||
assert credits_reader.display_summary() == "积分 12.35"
|
||||
|
||||
|
||||
def test_codex_reader_hides_reset_countdown_when_remaining_is_full() -> None:
|
||||
reader = get_quota_reader(
|
||||
"codex",
|
||||
{
|
||||
"codex": {
|
||||
"primary_used_percent": 0.0,
|
||||
"primary_reset_seconds": 266400,
|
||||
"secondary_used_percent": 0.0,
|
||||
"secondary_reset_seconds": 3600,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert reader.display_summary() == "周剩余 100.0% | 5H剩余 100.0%"
|
||||
|
||||
|
||||
def test_antigravity_reader_keeps_used_percent_fallbacks() -> None:
|
||||
reader = get_quota_reader(
|
||||
"antigravity",
|
||||
@@ -88,6 +104,20 @@ def test_extract_reset_seconds_uses_codex_weekly_reset_for_codex_provider() -> N
|
||||
assert extract_reset_seconds(key_obj) == pytest.approx(1800.0)
|
||||
|
||||
|
||||
def test_extract_reset_seconds_codex_weekly_full_quota_returns_none() -> None:
|
||||
key_obj = SimpleNamespace(
|
||||
provider_type="codex",
|
||||
upstream_metadata={
|
||||
"codex": {
|
||||
"primary_used_percent": 0.0,
|
||||
"primary_reset_seconds": 1800.0,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert extract_reset_seconds(key_obj) is None
|
||||
|
||||
|
||||
def test_extract_reset_seconds_prefers_codex_reset_at(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("src.services.provider.pool.dimensions._helpers.time.time", lambda: 1000.0)
|
||||
key_obj = SimpleNamespace(
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
@@ -124,44 +125,29 @@ async def test_candidate_cleanup_uses_dedicated_retention_and_batch_settings(
|
||||
assert batch_two.closed is True
|
||||
|
||||
|
||||
def test_cleanup_body_fields_loads_ids_then_processes_records_individually(
|
||||
def test_cleanup_body_fields_batches_records_with_single_commit(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
scheduler = MaintenanceScheduler()
|
||||
|
||||
class _IdBatchSession:
|
||||
def __init__(self, ids: list[str]) -> None:
|
||||
self.ids = ids
|
||||
self.closed = False
|
||||
self.query_obj = MagicMock()
|
||||
filtered = self.query_obj.filter.return_value
|
||||
filtered.filter.return_value.order_by.return_value.limit.return_value.all.return_value = [
|
||||
SimpleNamespace(id=value) for value in ids
|
||||
]
|
||||
|
||||
def query(self, *args): # type: ignore[no-untyped-def]
|
||||
self.query_args = args
|
||||
return self.query_obj
|
||||
|
||||
def rollback(self) -> None:
|
||||
raise AssertionError("rollback should not be called")
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
class _RecordSession:
|
||||
def __init__(self, record: SimpleNamespace) -> None:
|
||||
self.record = record
|
||||
class _BatchSession:
|
||||
def __init__(self, records: list[SimpleNamespace]) -> None:
|
||||
self.records = records
|
||||
self.closed = False
|
||||
self.committed = False
|
||||
self.executed = 0
|
||||
self.query_obj = MagicMock()
|
||||
self.query_obj.filter.return_value.first.return_value = record
|
||||
filtered = self.query_obj.filter.return_value
|
||||
filtered.filter.return_value.order_by.return_value.limit.return_value.all.return_value = (
|
||||
records
|
||||
)
|
||||
|
||||
def query(self, *args): # type: ignore[no-untyped-def]
|
||||
self.query_args = args
|
||||
return self.query_obj
|
||||
|
||||
def execute(self, _statement): # type: ignore[no-untyped-def]
|
||||
self.executed += 1
|
||||
return SimpleNamespace(rowcount=1)
|
||||
|
||||
def commit(self) -> None:
|
||||
@@ -173,27 +159,26 @@ def test_cleanup_body_fields_loads_ids_then_processes_records_individually(
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
batch_one = _IdBatchSession(["usage-1", "usage-2"])
|
||||
record_one = _RecordSession(
|
||||
SimpleNamespace(
|
||||
id="usage-1",
|
||||
request_body={"hello": "world"},
|
||||
response_body=None,
|
||||
provider_request_body=None,
|
||||
client_response_body=None,
|
||||
)
|
||||
batch_one = _BatchSession(
|
||||
[
|
||||
SimpleNamespace(
|
||||
id="usage-1",
|
||||
request_body={"hello": "world"},
|
||||
response_body=None,
|
||||
provider_request_body=None,
|
||||
client_response_body=None,
|
||||
),
|
||||
SimpleNamespace(
|
||||
id="usage-2",
|
||||
request_body=None,
|
||||
response_body={"ok": True},
|
||||
provider_request_body=None,
|
||||
client_response_body=None,
|
||||
),
|
||||
]
|
||||
)
|
||||
record_two = _RecordSession(
|
||||
SimpleNamespace(
|
||||
id="usage-2",
|
||||
request_body=None,
|
||||
response_body={"ok": True},
|
||||
provider_request_body=None,
|
||||
client_response_body=None,
|
||||
)
|
||||
)
|
||||
batch_two = _IdBatchSession([])
|
||||
sessions = iter([batch_one, record_one, record_two, batch_two])
|
||||
batch_two = _BatchSession([])
|
||||
sessions = iter([batch_one, batch_two])
|
||||
|
||||
monkeypatch.setattr(
|
||||
maintenance_scheduler_module,
|
||||
@@ -212,12 +197,173 @@ def test_cleanup_body_fields_loads_ids_then_processes_records_individually(
|
||||
)
|
||||
|
||||
assert compressed == 2
|
||||
assert batch_one.query_args == (maintenance_scheduler_module.Usage.id,)
|
||||
assert len(record_one.query_args) == 5
|
||||
assert len(record_two.query_args) == 5
|
||||
assert len(batch_one.query_args) == 5
|
||||
assert batch_one.executed == 2
|
||||
assert batch_one.committed is True
|
||||
assert batch_one.closed is True
|
||||
assert batch_two.closed is True
|
||||
assert record_one.committed is True
|
||||
assert record_two.committed is True
|
||||
assert record_one.closed is True
|
||||
assert record_two.closed is True
|
||||
|
||||
|
||||
def test_cleanup_header_fields_clears_client_response_headers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
scheduler = MaintenanceScheduler()
|
||||
|
||||
class _BatchSession:
|
||||
def __init__(self, ids: list[str]) -> None:
|
||||
self.ids = ids
|
||||
self.closed = False
|
||||
self.committed = False
|
||||
self.query_obj = MagicMock()
|
||||
self.filtered_by_time = MagicMock()
|
||||
self.filtered_by_headers = MagicMock()
|
||||
self.query_obj.filter.return_value = self.filtered_by_time
|
||||
self.filtered_by_time.filter.return_value = self.filtered_by_headers
|
||||
self.filtered_by_headers.order_by.return_value.limit.return_value.all.return_value = [
|
||||
SimpleNamespace(id=value) for value in ids
|
||||
]
|
||||
self.executed_statements: list[str] = []
|
||||
|
||||
def query(self, *args): # type: ignore[no-untyped-def]
|
||||
self.query_args = args
|
||||
return self.query_obj
|
||||
|
||||
def execute(self, statement): # type: ignore[no-untyped-def]
|
||||
self.executed_statements.append(str(statement))
|
||||
return SimpleNamespace(rowcount=len(self.ids))
|
||||
|
||||
def commit(self) -> None:
|
||||
self.committed = True
|
||||
|
||||
def rollback(self) -> None:
|
||||
raise AssertionError("rollback should not be called")
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
batch_one = _BatchSession(["usage-1"])
|
||||
batch_two = _BatchSession([])
|
||||
sessions = iter([batch_one, batch_two])
|
||||
|
||||
monkeypatch.setattr(
|
||||
maintenance_scheduler_module,
|
||||
"create_session",
|
||||
lambda: next(sessions),
|
||||
)
|
||||
|
||||
cleaned = scheduler._cleanup_header_fields(
|
||||
cutoff_time=SimpleNamespace(), # type: ignore[arg-type]
|
||||
batch_size=1000,
|
||||
)
|
||||
|
||||
header_filter = str(batch_one.filtered_by_time.filter.call_args.args[0])
|
||||
|
||||
assert cleaned == 1
|
||||
assert batch_one.query_args == (maintenance_scheduler_module.Usage.id,)
|
||||
assert "client_response_headers" in header_filter
|
||||
assert "client_response_headers" in batch_one.executed_statements[0]
|
||||
assert batch_one.committed is True
|
||||
assert batch_one.closed is True
|
||||
assert batch_two.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_perform_cleanup_deletes_first_and_uses_non_overlapping_windows(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
scheduler = MaintenanceScheduler()
|
||||
fixed_now = datetime(2026, 3, 18, 3, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
class _FakeDateTime(datetime):
|
||||
@classmethod
|
||||
def now(cls, tz=None): # type: ignore[override]
|
||||
if tz is None:
|
||||
return fixed_now.replace(tzinfo=None)
|
||||
return fixed_now.astimezone(tz)
|
||||
|
||||
class _FakeLoop:
|
||||
async def run_in_executor(self, _executor, func): # type: ignore[no-untyped-def]
|
||||
return func()
|
||||
|
||||
class _ConfigSession:
|
||||
def close(self) -> None:
|
||||
return None
|
||||
|
||||
calls: list[tuple[str, datetime, int, datetime | None]] = []
|
||||
|
||||
def _record(name: str, count: int):
|
||||
def _inner(
|
||||
cutoff_time: datetime,
|
||||
batch_size: int,
|
||||
*,
|
||||
newer_than: datetime | None = None,
|
||||
) -> int:
|
||||
calls.append((name, cutoff_time, batch_size, newer_than))
|
||||
return count
|
||||
|
||||
return _inner
|
||||
|
||||
config_values = {
|
||||
"enable_auto_cleanup": True,
|
||||
"detail_log_retention_days": 7,
|
||||
"compressed_log_retention_days": 30,
|
||||
"header_retention_days": 90,
|
||||
"log_retention_days": 365,
|
||||
"cleanup_batch_size": 123,
|
||||
"auto_delete_expired_keys": False,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(maintenance_scheduler_module, "datetime", _FakeDateTime)
|
||||
monkeypatch.setattr(
|
||||
maintenance_scheduler_module.asyncio, "get_running_loop", lambda: _FakeLoop()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
maintenance_scheduler_module,
|
||||
"create_session",
|
||||
lambda: _ConfigSession(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
maintenance_scheduler_module.SystemConfigService,
|
||||
"get_config",
|
||||
lambda _db, key, default=None: config_values.get(key, default),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
scheduler,
|
||||
"_delete_old_records",
|
||||
lambda cutoff_time, batch_size: calls.append(("delete", cutoff_time, batch_size, None))
|
||||
or 5,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
scheduler,
|
||||
"_cleanup_header_fields",
|
||||
_record("header", 4),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
scheduler,
|
||||
"_cleanup_stale_body_fields",
|
||||
_record("body", 3),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
scheduler,
|
||||
"_cleanup_body_fields",
|
||||
_record("compress", 2),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
maintenance_scheduler_module.ApiKeyService,
|
||||
"cleanup_expired_keys",
|
||||
lambda _db, auto_delete=False: 0,
|
||||
)
|
||||
|
||||
await scheduler._perform_cleanup()
|
||||
|
||||
detail_cutoff = fixed_now - timedelta(days=7)
|
||||
compressed_cutoff = fixed_now - timedelta(days=30)
|
||||
header_cutoff = fixed_now - timedelta(days=90)
|
||||
log_cutoff = fixed_now - timedelta(days=365)
|
||||
|
||||
assert calls == [
|
||||
("delete", log_cutoff, 123, None),
|
||||
("header", header_cutoff, 123, log_cutoff),
|
||||
("body", compressed_cutoff, 123, log_cutoff),
|
||||
("compress", detail_cutoff, 123, compressed_cutoff),
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user