mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(api-keys): 支持独立密钥额度重置(手动+定时自动)
- 新增手动重置接口 PATCH /api/admin/api-keys/{id}/reset-usage
- 前端 ApiKeys 页面增加重置按钮,支持确认后归零已使用额度
- 新增独立密钥额度定时自动重置任务,支持配置周期和执行时间
- 支持 all/selected 两种重置模式,selected 模式可指定密钥
- 移除已废弃的 CleanupScheduler 兼容别名
This commit is contained in:
@@ -292,6 +292,26 @@ async def add_balance_to_key(
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.patch("/{key_id}/reset-usage")
|
||||
async def reset_api_key_usage(key_id: str, request: Request, db: Session = Depends(get_db)) -> Any:
|
||||
"""
|
||||
重置独立余额 API Key 的已使用额度
|
||||
|
||||
将 balance_used_usd 重置为 0,不改变 current_balance_usd。
|
||||
|
||||
**路径参数**:
|
||||
- `key_id`: API Key ID
|
||||
|
||||
**返回字段**:
|
||||
- `id`: API Key ID
|
||||
- `current_balance_usd`: 当前余额
|
||||
- `balance_used_usd`: 已使用余额(重置后为 0)
|
||||
- `message`: 提示信息
|
||||
"""
|
||||
adapter = AdminResetKeyUsageAdapter(key_id=key_id)
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.get("/{key_id}")
|
||||
async def get_api_key_detail(
|
||||
key_id: str,
|
||||
@@ -663,6 +683,48 @@ class AdminAddBalanceAdapter(AdminApiAdapter):
|
||||
}
|
||||
|
||||
|
||||
class AdminResetKeyUsageAdapter(AdminApiAdapter):
|
||||
"""重置独立余额Key的已使用额度"""
|
||||
|
||||
def __init__(self, key_id: str):
|
||||
self.key_id = key_id
|
||||
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
db = context.db
|
||||
api_key = db.query(ApiKey).filter(ApiKey.id == self.key_id).first()
|
||||
if not api_key:
|
||||
raise NotFoundException("API密钥不存在", "api_key")
|
||||
|
||||
if not api_key.is_standalone:
|
||||
raise InvalidRequestException("只能重置独立余额Key的使用额度")
|
||||
|
||||
previous_used = float(api_key.balance_used_usd or 0)
|
||||
api_key.balance_used_usd = 0.0
|
||||
api_key.updated_at = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
db.refresh(api_key)
|
||||
|
||||
logger.info(
|
||||
f"管理员重置独立余额Key使用额度: Key ID {self.key_id}, "
|
||||
f"重置前已使用 ${previous_used:.4f}"
|
||||
)
|
||||
|
||||
context.add_audit_metadata(
|
||||
action="reset_key_usage",
|
||||
key_id=self.key_id,
|
||||
current_balance_usd=api_key.current_balance_usd,
|
||||
previous_balance_used_usd=previous_used,
|
||||
)
|
||||
|
||||
return {
|
||||
"id": api_key.id,
|
||||
"name": api_key.name,
|
||||
"current_balance_usd": api_key.current_balance_usd,
|
||||
"balance_used_usd": 0.0,
|
||||
"message": "使用额度已重置",
|
||||
}
|
||||
|
||||
|
||||
class AdminGetFullKeyAdapter(AdminApiAdapter):
|
||||
"""获取完整的API密钥"""
|
||||
|
||||
|
||||
@@ -684,6 +684,16 @@ class AdminSetSystemConfigAdapter(AdminApiAdapter):
|
||||
except Exception as e:
|
||||
logger.warning(f"更新用户配额重置任务时间失败: {e}")
|
||||
|
||||
# 如果更新的是独立密钥额度重置任务时间,动态更新调度器
|
||||
if self.key == "standalone_key_quota_reset_time" and value:
|
||||
try:
|
||||
from src.services.system.maintenance_scheduler import get_maintenance_scheduler
|
||||
|
||||
scheduler = get_maintenance_scheduler()
|
||||
scheduler.update_standalone_key_quota_reset_time(value)
|
||||
except Exception as e:
|
||||
logger.warning(f"更新独立密钥额度重置任务时间失败: {e}")
|
||||
|
||||
# 如果更新的是调度模式或优先级模式,立即更新当前 Worker 的 Scheduler 单例
|
||||
if self.key in ("scheduling_mode", "provider_priority_mode"):
|
||||
try:
|
||||
@@ -2444,9 +2454,7 @@ class AdminPurgeUsersAdapter(AdminApiAdapter):
|
||||
|
||||
db = context.db
|
||||
|
||||
user_ids = [
|
||||
uid for (uid,) in db.query(User.id).filter(User.role != UserRole.ADMIN).all()
|
||||
]
|
||||
user_ids = [uid for (uid,) in db.query(User.id).filter(User.role != UserRole.ADMIN).all()]
|
||||
users_count = len(user_ids)
|
||||
|
||||
if user_ids:
|
||||
|
||||
@@ -7,7 +7,6 @@
|
||||
from src.services.system.announcement import AnnouncementService
|
||||
from src.services.system.audit import AuditService
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.system.maintenance_scheduler import CleanupScheduler # 兼容旧名称
|
||||
from src.services.system.maintenance_scheduler import (
|
||||
MaintenanceScheduler,
|
||||
get_maintenance_scheduler,
|
||||
@@ -20,7 +19,6 @@ __all__ = [
|
||||
"AuditService",
|
||||
"AnnouncementService",
|
||||
"MaintenanceScheduler",
|
||||
"CleanupScheduler", # 兼容旧名称
|
||||
"get_maintenance_scheduler",
|
||||
"SyncStatsService",
|
||||
"TaskScheduler",
|
||||
|
||||
@@ -155,6 +155,26 @@ class SystemConfigService:
|
||||
"value": 1,
|
||||
"description": "用户配额重置周期(天数)",
|
||||
},
|
||||
"enable_standalone_key_quota_reset": {
|
||||
"value": False,
|
||||
"description": "是否启用独立密钥额度自动重置任务(按配置时间触发,按周期执行)",
|
||||
},
|
||||
"standalone_key_quota_reset_time": {
|
||||
"value": "05:00",
|
||||
"description": "独立密钥额度自动重置执行时间(HH:MM 格式,24小时制)",
|
||||
},
|
||||
"standalone_key_quota_reset_interval_days": {
|
||||
"value": 1,
|
||||
"description": "独立密钥额度重置周期(天数)",
|
||||
},
|
||||
"standalone_key_quota_reset_mode": {
|
||||
"value": "all",
|
||||
"description": "独立密钥额度重置模式:all(全部独立密钥) 或 selected(指定密钥)",
|
||||
},
|
||||
"standalone_key_quota_reset_key_ids": {
|
||||
"value": [],
|
||||
"description": "独立密钥额度重置指定的密钥 ID 列表(仅 mode=selected 时生效)",
|
||||
},
|
||||
"provider_priority_mode": {
|
||||
"value": "provider",
|
||||
"description": "优先级策略:provider(提供商优先模式) 或 global_key(全局Key优先模式)",
|
||||
|
||||
@@ -24,7 +24,7 @@ from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import AuditLog, Provider, Usage
|
||||
from src.models.database import ApiKey, AuditLog, Provider, Usage
|
||||
from src.services.provider_ops.service import ProviderOpsService
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.system.scheduler import get_scheduler
|
||||
@@ -40,6 +40,8 @@ class MaintenanceScheduler:
|
||||
CHECKIN_JOB_ID = "provider_checkin"
|
||||
# 用户配额重置任务的 job_id
|
||||
USER_QUOTA_RESET_JOB_ID = "user_quota_reset"
|
||||
# 独立密钥额度重置任务的 job_id
|
||||
STANDALONE_KEY_QUOTA_RESET_JOB_ID = "standalone_key_quota_reset"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.running = False
|
||||
@@ -161,6 +163,33 @@ class MaintenanceScheduler:
|
||||
|
||||
return success
|
||||
|
||||
def _get_standalone_key_quota_reset_time(self) -> tuple[int, int]:
|
||||
"""获取独立密钥额度重置任务的执行时间"""
|
||||
db = create_session()
|
||||
try:
|
||||
time_str = SystemConfigService.get_config(
|
||||
db, "standalone_key_quota_reset_time", "05:00"
|
||||
)
|
||||
return self._parse_user_quota_reset_time_string(time_str)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def update_standalone_key_quota_reset_time(self, time_str: str) -> bool:
|
||||
"""更新独立密钥额度重置任务的执行时间"""
|
||||
hour, minute = self._parse_user_quota_reset_time_string(time_str)
|
||||
|
||||
scheduler = get_scheduler()
|
||||
success = scheduler.reschedule_cron_job(
|
||||
self.STANDALONE_KEY_QUOTA_RESET_JOB_ID,
|
||||
hour=hour,
|
||||
minute=minute,
|
||||
)
|
||||
|
||||
if success:
|
||||
logger.info(f"独立密钥额度重置任务时间已更新为: {hour:02d}:{minute:02d}")
|
||||
|
||||
return success
|
||||
|
||||
def get_checkin_job_info(self) -> dict | None:
|
||||
"""获取签到任务的信息
|
||||
|
||||
@@ -279,6 +308,16 @@ class MaintenanceScheduler:
|
||||
name="用户配额自动重置",
|
||||
)
|
||||
|
||||
# 独立密钥额度重置任务 - 根据配置时间执行(按周期配置决定是否执行)
|
||||
sk_reset_hour, sk_reset_minute = self._get_standalone_key_quota_reset_time()
|
||||
scheduler.add_cron_job(
|
||||
self._scheduled_standalone_key_quota_reset,
|
||||
hour=sk_reset_hour,
|
||||
minute=sk_reset_minute,
|
||||
job_id=self.STANDALONE_KEY_QUOTA_RESET_JOB_ID,
|
||||
name="独立密钥额度自动重置",
|
||||
)
|
||||
|
||||
# 启动时执行一次初始化任务
|
||||
asyncio.create_task(self._run_startup_tasks())
|
||||
|
||||
@@ -371,6 +410,10 @@ class MaintenanceScheduler:
|
||||
"""用户配额重置任务(定时调用)"""
|
||||
await self._perform_user_quota_reset()
|
||||
|
||||
async def _scheduled_standalone_key_quota_reset(self) -> None:
|
||||
"""独立密钥额度重置任务(定时调用)"""
|
||||
await self._perform_standalone_key_quota_reset()
|
||||
|
||||
# ========== 实际任务实现 ==========
|
||||
|
||||
async def _perform_stats_aggregation(self, backfill: bool = False) -> None:
|
||||
@@ -876,6 +919,125 @@ class MaintenanceScheduler:
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def _perform_standalone_key_quota_reset(self) -> None:
|
||||
"""执行独立密钥额度自动重置任务
|
||||
|
||||
适用范围:
|
||||
- is_standalone=True 的密钥
|
||||
- current_balance_usd != NULL(有限额的密钥)
|
||||
- 支持 all(全部)和 selected(指定密钥)两种模式
|
||||
"""
|
||||
db = create_session()
|
||||
try:
|
||||
if not SystemConfigService.get_config(db, "enable_standalone_key_quota_reset", False):
|
||||
logger.info("独立密钥额度自动重置已禁用,跳过任务")
|
||||
return
|
||||
|
||||
# 重置周期
|
||||
interval_value = SystemConfigService.get_config(
|
||||
db, "standalone_key_quota_reset_interval_days", 1
|
||||
)
|
||||
try:
|
||||
interval_days = int(interval_value)
|
||||
except Exception:
|
||||
interval_days = 1
|
||||
if interval_days < 1:
|
||||
interval_days = 1
|
||||
|
||||
# 滚动计算
|
||||
last_reset_at = SystemConfigService.get_config(db, "standalone_key_quota_last_reset_at")
|
||||
|
||||
should_run = True
|
||||
if last_reset_at:
|
||||
last_dt: datetime | None = None
|
||||
try:
|
||||
if isinstance(last_reset_at, str):
|
||||
last_dt = datetime.fromisoformat(last_reset_at)
|
||||
except Exception:
|
||||
last_dt = None
|
||||
|
||||
if last_dt is None:
|
||||
logger.warning("standalone_key_quota_last_reset_at 格式无效,视为需要执行一次")
|
||||
else:
|
||||
if last_dt.tzinfo is None:
|
||||
last_dt = last_dt.replace(tzinfo=timezone.utc)
|
||||
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from src.services.system.scheduler import APP_TIMEZONE
|
||||
|
||||
tz = ZoneInfo(APP_TIMEZONE)
|
||||
now_local = datetime.now(tz)
|
||||
last_local_date = last_dt.astimezone(tz).date()
|
||||
days_since_reset = (now_local.date() - last_local_date).days
|
||||
|
||||
if days_since_reset < 0:
|
||||
logger.warning("standalone_key_quota_last_reset_at 在未来,跳过本次重置")
|
||||
should_run = False
|
||||
elif days_since_reset < interval_days:
|
||||
logger.info(
|
||||
f"独立密钥额度自动重置未到周期,跳过任务"
|
||||
f"({days_since_reset}/{interval_days}天)"
|
||||
)
|
||||
should_run = False
|
||||
|
||||
if not should_run:
|
||||
return
|
||||
|
||||
# 确定重置范围
|
||||
reset_mode = SystemConfigService.get_config(
|
||||
db, "standalone_key_quota_reset_mode", "all"
|
||||
)
|
||||
|
||||
now_utc = datetime.now(timezone.utc)
|
||||
base_filter = [
|
||||
ApiKey.is_standalone.is_(True),
|
||||
ApiKey.current_balance_usd.isnot(None),
|
||||
]
|
||||
|
||||
if reset_mode == "selected":
|
||||
key_ids = SystemConfigService.get_config(
|
||||
db, "standalone_key_quota_reset_key_ids", []
|
||||
)
|
||||
if not key_ids:
|
||||
logger.info("独立密钥额度重置模式为 selected 但未选择任何密钥,跳过")
|
||||
return
|
||||
base_filter.append(ApiKey.id.in_(key_ids))
|
||||
|
||||
reset_count = (
|
||||
db.query(ApiKey)
|
||||
.filter(*base_filter)
|
||||
.update(
|
||||
{
|
||||
ApiKey.balance_used_usd: 0.0,
|
||||
ApiKey.updated_at: now_utc,
|
||||
},
|
||||
synchronize_session=False,
|
||||
)
|
||||
)
|
||||
db.commit()
|
||||
|
||||
SystemConfigService.set_config(
|
||||
db,
|
||||
"standalone_key_quota_last_reset_at",
|
||||
now_utc.isoformat(),
|
||||
"独立密钥额度自动重置的上次执行时间(UTC,内部使用)",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"独立密钥额度自动重置完成: mode={reset_mode}, "
|
||||
f"interval_days={interval_days}, 重置密钥数={reset_count}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"独立密钥额度自动重置任务执行失败: {e}")
|
||||
try:
|
||||
db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def _perform_cleanup(self) -> None:
|
||||
"""执行清理任务"""
|
||||
db = create_session()
|
||||
@@ -1276,12 +1438,3 @@ def get_maintenance_scheduler() -> MaintenanceScheduler:
|
||||
if _maintenance_scheduler is None:
|
||||
_maintenance_scheduler = MaintenanceScheduler()
|
||||
return _maintenance_scheduler
|
||||
|
||||
|
||||
# 兼容旧名称(deprecated)
|
||||
def get_cleanup_scheduler() -> MaintenanceScheduler:
|
||||
"""获取维护调度器单例(已废弃,请使用 get_maintenance_scheduler)"""
|
||||
return get_maintenance_scheduler()
|
||||
|
||||
|
||||
CleanupScheduler = MaintenanceScheduler # 兼容旧名称
|
||||
|
||||
Reference in New Issue
Block a user