Download services/auth_service.py from zlcy/study: direct link, hf CLI and curl.
- Browser
- Download file 142 kB
-
https://huggingface.co/spaces/zlcy/study/resolve/main/services/auth_service.py
- Command line
-
hf download hf://spaces/zlcy/study/services/auth_service.py
-
curl -L -o auth_service.py https://huggingface.co/spaces/zlcy/study/resolve/main/services/auth_service.py
142 kB
| from __future__ import annotations | |
| import hashlib | |
| import hmac | |
| import secrets | |
| import time | |
| import uuid | |
| from datetime import datetime, timedelta, timezone | |
| from decimal import Decimal, InvalidOperation, ROUND_HALF_UP | |
| from threading import RLock, Thread | |
| from typing import Callable, Literal, TypeVar | |
| try: | |
| from zoneinfo import ZoneInfo | |
| except ImportError: # pragma: no cover - Python 3.9+ normally has zoneinfo. | |
| ZoneInfo = None # type: ignore[assignment] | |
| from services.config import config | |
| from services.storage.base import StorageBackend | |
| AuthRole = Literal["admin", "user", "normal"] | |
| ResultT = TypeVar("ResultT") | |
| THIRD_PARTY_IMAGE_GROUP_WALLET_MODEL_ID = "__group__" | |
| DAILY_GIFT_ACCOUNT_POOL_MODEL_ID = "__account_pool__" | |
| def _now_iso() -> str: | |
| return datetime.now(timezone.utc).isoformat() | |
| def _today_key() -> str: | |
| tz = ZoneInfo("Asia/Shanghai") if ZoneInfo else timezone(timedelta(hours=8)) | |
| return datetime.now(tz).strftime("%Y-%m-%d") | |
| def _china_now() -> datetime: | |
| tz = ZoneInfo("Asia/Shanghai") if ZoneInfo else timezone(timedelta(hours=8)) | |
| return datetime.now(tz) | |
| def _period_keys(now: datetime | None = None) -> tuple[str, str, str]: | |
| current = now or _china_now() | |
| date_key = current.strftime("%Y-%m-%d") | |
| iso_year, iso_week, _ = current.isocalendar() | |
| return date_key, f"{iso_year}-W{iso_week:02d}", current.strftime("%Y-%m") | |
| def _hash_key(value: str) -> str: | |
| return hashlib.sha256(value.encode("utf-8")).hexdigest() | |
| def _invite_code() -> str: | |
| return secrets.token_urlsafe(6).replace("-", "").replace("_", "").upper()[:8] | |
| def _to_nonnegative_int(value: object) -> int: | |
| try: | |
| return max(0, int(value or 0)) | |
| except (TypeError, ValueError): | |
| return 0 | |
| CREDIT_QUANTUM = Decimal("0.01") | |
| def _credit_decimal(value: object) -> Decimal: | |
| try: | |
| parsed = Decimal(str(value or 0)) | |
| except (InvalidOperation, TypeError, ValueError): | |
| return Decimal("0.00") | |
| return max(Decimal("0"), parsed).quantize(CREDIT_QUANTUM, rounding=ROUND_HALF_UP) | |
| def _credit_number(value: object) -> int | float: | |
| parsed = _credit_decimal(value) | |
| return int(parsed) if parsed == parsed.to_integral_value() else float(parsed) | |
| def _parse_iso_datetime(value: object) -> datetime | None: | |
| raw = str(value or "").strip() | |
| if not raw: | |
| return None | |
| try: | |
| parsed = datetime.fromisoformat(raw.replace("Z", "+00:00")) | |
| except ValueError: | |
| return None | |
| return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc) | |
| def _is_future_datetime(value: object) -> bool: | |
| parsed = _parse_iso_datetime(value) | |
| return parsed is not None and parsed > datetime.now(timezone.utc) | |
| def _normalize_credit_grants(value: object) -> list[dict[str, object]]: | |
| source = value if isinstance(value, list) else [] | |
| grants: list[dict[str, object]] = [] | |
| for raw in source: | |
| if not isinstance(raw, dict): | |
| continue | |
| grant_id = str(raw.get("id") or "").strip() or uuid.uuid4().hex[:12] | |
| expires_at = str(raw.get("expires_at") or "").strip() | |
| if not expires_at: | |
| continue | |
| grants.append( | |
| { | |
| "id": grant_id, | |
| "remaining": _credit_number(raw.get("remaining")), | |
| "reserved": _credit_number(raw.get("reserved")), | |
| "expires_at": expires_at, | |
| "created_at": str(raw.get("created_at") or "").strip() or _now_iso(), | |
| "source": str(raw.get("source") or "").strip(), | |
| "code_id": str(raw.get("code_id") or "").strip(), | |
| } | |
| ) | |
| grants.sort(key=lambda item: str(item.get("expires_at") or "")) | |
| return grants | |
| def _active_credit_grants(item: dict[str, object]) -> list[dict[str, object]]: | |
| return [grant for grant in _normalize_credit_grants(item.get("credit_grants")) if _is_future_datetime(grant.get("expires_at"))] | |
| def _expiring_credit_remaining(item: dict[str, object]) -> int | float: | |
| return _credit_number(sum((_credit_decimal(grant.get("remaining")) for grant in _active_credit_grants(item)), Decimal("0"))) | |
| def _expiring_credit_reserved(item: dict[str, object]) -> int | float: | |
| return _credit_number(sum((_credit_decimal(grant.get("reserved")) for grant in _normalize_credit_grants(item.get("credit_grants"))), Decimal("0"))) | |
| CHECKIN_CREDIT_SOURCE = "checkin" | |
| def _checkin_credit_state(item: dict[str, object]) -> tuple[Decimal, str]: | |
| """返回 (该用户还没用掉的签到积分, 最近到期时间)。 | |
| 签到赠送积分有两种形态,取决于后台「签到积分可用天数」: | |
| - 永久形态(可用天数 = 0,默认):直接进基础积分池,由 `checkin_credit_remaining` 账本跟踪; | |
| - 限时形态(可用天数 > 0):发放成 `source="checkin"` 的限时券,剩余量记在券上,到期即作废。 | |
| 两者相加才是「未使用的签到积分」;`expires_at` 取未过期签到券里最早的那张(没有则为空串)。 | |
| """ | |
| base_balance = _credit_decimal(item.get("credits")) + _credit_decimal(item.get("reserved_credits")) | |
| total = min(_credit_decimal(item.get("checkin_credit_remaining")), base_balance) | |
| expires_at = "" | |
| for grant in _active_credit_grants(item): | |
| if str(grant.get("source") or "").strip() != CHECKIN_CREDIT_SOURCE: | |
| continue | |
| remaining = _credit_decimal(grant.get("remaining")) | |
| total += remaining | |
| # 剩余为 0 的券(已被扣光但还没清理)不该参与「最近到期时间」,否则会把它当成最近的一笔。 | |
| if remaining <= 0: | |
| continue | |
| candidate = str(grant.get("expires_at") or "").strip() | |
| if candidate and (not expires_at or candidate < expires_at): | |
| expires_at = candidate | |
| return total, expires_at | |
| def _checkin_credit_batches(item: dict[str, object]) -> list[dict[str, object]]: | |
| """把「未过期的签到限时券」按到期时刻聚合成 [{credits, expires_at}],按到期时间升序。 | |
| 签到积分逐笔到期(每笔自发放起 N 天后失效),所以这里会天然出现多批;同一时刻到期的 | |
| 券(例如后台保存设置触发重算后对齐的那批)会合并成一条。 | |
| 永久形态(可用天数 = 0)没有到期时间,不产生批次 —— 前端据此隐藏到期那几行。 | |
| """ | |
| grouped: dict[str, Decimal] = {} | |
| for grant in _active_credit_grants(item): | |
| if str(grant.get("source") or "").strip() != CHECKIN_CREDIT_SOURCE: | |
| continue | |
| remaining = _credit_decimal(grant.get("remaining")) | |
| if remaining <= 0: | |
| continue | |
| key = str(grant.get("expires_at") or "").strip() | |
| grouped[key] = grouped.get(key, Decimal("0")) + remaining | |
| return [ | |
| {"credits": _credit_number(total), "expires_at": expires_at} | |
| for expires_at, total in sorted(grouped.items(), key=lambda pair: pair[0]) | |
| ] | |
| PACKAGE_UNLIMITED_SENTINEL = 10**9 | |
| def _normalize_package_grants(value: object) -> list[dict[str, object]]: | |
| source = value if isinstance(value, list) else [] | |
| grants: list[dict[str, object]] = [] | |
| for raw in source: | |
| if not isinstance(raw, dict): | |
| continue | |
| grant_id = str(raw.get("id") or "").strip() or uuid.uuid4().hex[:12] | |
| expires_at = str(raw.get("expires_at") or "").strip() | |
| if not expires_at: | |
| continue | |
| grants.append( | |
| { | |
| "id": grant_id, | |
| "name": str(raw.get("name") or "套餐").strip() or "套餐", | |
| "created_at": str(raw.get("created_at") or "").strip() or _now_iso(), | |
| "starts_at": str(raw.get("starts_at") or "").strip(), | |
| "expires_at": expires_at, | |
| "source": str(raw.get("source") or "").strip(), | |
| "code_id": str(raw.get("code_id") or "").strip(), | |
| "daily_limit": _to_nonnegative_int(raw.get("daily_limit")), | |
| "weekly_limit": _to_nonnegative_int(raw.get("weekly_limit")), | |
| "monthly_limit": _to_nonnegative_int(raw.get("monthly_limit")), | |
| "total_limit": _to_nonnegative_int(raw.get("total_limit")), | |
| "used_total": _credit_number(raw.get("used_total")), | |
| "reserved_total": _credit_number(raw.get("reserved_total")), | |
| "daily_key": str(raw.get("daily_key") or "").strip(), | |
| "used_daily": _credit_number(raw.get("used_daily")), | |
| "reserved_daily": _credit_number(raw.get("reserved_daily")), | |
| "weekly_key": str(raw.get("weekly_key") or "").strip(), | |
| "used_weekly": _credit_number(raw.get("used_weekly")), | |
| "reserved_weekly": _credit_number(raw.get("reserved_weekly")), | |
| "monthly_key": str(raw.get("monthly_key") or "").strip(), | |
| "used_monthly": _credit_number(raw.get("used_monthly")), | |
| "reserved_monthly": _credit_number(raw.get("reserved_monthly")), | |
| } | |
| ) | |
| grants.sort(key=lambda item: str(item.get("expires_at") or "")) | |
| return grants | |
| def _normalize_string_list(value: object) -> list[str]: | |
| source = value if isinstance(value, list) else [] | |
| result: list[str] = [] | |
| seen: set[str] = set() | |
| for item in source: | |
| text = str(item or "").strip() | |
| if not text or text in seen: | |
| continue | |
| seen.add(text) | |
| result.append(text) | |
| return result | |
| def _normalize_third_party_image_wallets(value: object) -> list[dict[str, object]]: | |
| source = value if isinstance(value, list) else [] | |
| wallets_by_key: dict[str, dict[str, object]] = {} | |
| for raw in source: | |
| if not isinstance(raw, dict): | |
| continue | |
| group_id = str(raw.get("group_id") or "").strip() | |
| model_id = str(raw.get("model_id") or "").strip() or THIRD_PARTY_IMAGE_GROUP_WALLET_MODEL_ID | |
| if not group_id: | |
| continue | |
| wallet_id = f"{group_id}:{model_id}" | |
| wallet = wallets_by_key.get(wallet_id) | |
| if wallet is None: | |
| wallet = { | |
| "id": wallet_id, | |
| "group_id": group_id, | |
| "model_id": model_id, | |
| "remaining": _credit_number(raw.get("remaining")), | |
| "reserved": _credit_number(raw.get("reserved")), | |
| "used": _credit_number(raw.get("used")), | |
| "total_recharged": _credit_number(raw.get("total_recharged")), | |
| "updated_at": str(raw.get("updated_at") or "").strip(), | |
| } | |
| wallets_by_key[wallet_id] = wallet | |
| else: | |
| wallet["remaining"] = _credit_number(_credit_decimal(wallet.get("remaining")) + _credit_decimal(raw.get("remaining"))) | |
| wallet["reserved"] = _credit_number(_credit_decimal(wallet.get("reserved")) + _credit_decimal(raw.get("reserved"))) | |
| wallet["used"] = _credit_number(_credit_decimal(wallet.get("used")) + _credit_decimal(raw.get("used"))) | |
| wallet["total_recharged"] = _credit_number(_credit_decimal(wallet.get("total_recharged")) + _credit_decimal(raw.get("total_recharged"))) | |
| updated_at = str(raw.get("updated_at") or "").strip() | |
| if updated_at > str(wallet.get("updated_at") or ""): | |
| wallet["updated_at"] = updated_at | |
| wallets = list(wallets_by_key.values()) | |
| wallets.sort(key=lambda item: (str(item.get("group_id") or ""), str(item.get("model_id") or ""))) | |
| return wallets | |
| def _normalize_image_model_daily_limits(value: object) -> list[dict[str, object]]: | |
| if isinstance(value, dict): | |
| source = [ | |
| {"model_id": model_id, **entry} if isinstance(entry, dict) else {"model_id": model_id, "daily_limit": entry} | |
| for model_id, entry in value.items() | |
| ] | |
| else: | |
| source = value if isinstance(value, list) else [] | |
| limits_by_model: dict[str, dict[str, object]] = {} | |
| for raw in source: | |
| if not isinstance(raw, dict): | |
| continue | |
| model_id = str(raw.get("model_id") or raw.get("model") or raw.get("id") or "").strip() | |
| if not model_id: | |
| continue | |
| previous = limits_by_model.get(model_id) | |
| used = _to_nonnegative_int(raw.get("used")) | |
| reserved = _to_nonnegative_int(raw.get("reserved")) | |
| if previous is not None and str(previous.get("date") or "") == str(raw.get("date") or raw.get("daily_key") or "").strip(): | |
| used += _to_nonnegative_int(previous.get("used")) | |
| reserved += _to_nonnegative_int(previous.get("reserved")) | |
| limits_by_model[model_id] = { | |
| "model_id": model_id, | |
| "daily_limit": _to_nonnegative_int(raw.get("daily_limit") if "daily_limit" in raw else raw.get("limit")), | |
| "date": str(raw.get("date") or raw.get("daily_key") or "").strip(), | |
| "used": used, | |
| "reserved": reserved, | |
| "updated_at": str(raw.get("updated_at") or "").strip(), | |
| } | |
| limits = list(limits_by_model.values()) | |
| limits.sort(key=lambda item: str(item.get("model_id") or "")) | |
| return limits | |
| def _reset_image_model_daily_limit(entry: dict[str, object]) -> dict[str, object]: | |
| today = _today_key() | |
| next_entry = dict(entry) | |
| if str(next_entry.get("date") or "").strip() != today and _to_nonnegative_int(next_entry.get("reserved")) <= 0: | |
| next_entry["date"] = today | |
| next_entry["used"] = 0 | |
| next_entry["reserved"] = 0 | |
| return next_entry | |
| def _image_model_daily_limit_public(entry: dict[str, object]) -> dict[str, object]: | |
| normalized = _reset_image_model_daily_limit(entry) | |
| daily_limit = _to_nonnegative_int(normalized.get("daily_limit")) | |
| used = _to_nonnegative_int(normalized.get("used")) | |
| reserved = _to_nonnegative_int(normalized.get("reserved")) | |
| return { | |
| "model_id": str(normalized.get("model_id") or "").strip(), | |
| "daily_limit": daily_limit, | |
| "used": used, | |
| "reserved": reserved, | |
| "remaining": max(0, daily_limit - used - reserved), | |
| "date": str(normalized.get("date") or "").strip() or _today_key(), | |
| "updated_at": normalized.get("updated_at"), | |
| } | |
| def _third_party_wallet_key(group_id: str, model_id: str) -> str: | |
| group_id = str(group_id or "").strip() | |
| model_id = str(model_id or "").strip() | |
| return f"{group_id}:{model_id}" if group_id and model_id else "" | |
| def _is_matching_third_party_wallet(wallet: dict[str, object], group_id: str, model_id: str) -> bool: | |
| group_id = str(group_id or "").strip() | |
| model_id = str(model_id or "").strip() | |
| wallet_group_id = str(wallet.get("group_id") or "").strip() | |
| wallet_model_id = str(wallet.get("model_id") or "").strip() | |
| wallet_id = str(wallet.get("id") or "").strip() | |
| if wallet_id == _third_party_wallet_key(group_id, model_id): | |
| return True | |
| return wallet_group_id == group_id and wallet_model_id == model_id | |
| def _find_third_party_wallet_for_reserve( | |
| wallets: list[dict[str, object]], | |
| group_id: str, | |
| model_id: str, | |
| amount: int | float, | |
| ) -> dict[str, object] | None: | |
| matches = [wallet for wallet in wallets if _is_matching_third_party_wallet(wallet, group_id, model_id)] | |
| if not matches: | |
| return None | |
| for wallet in matches: | |
| if _credit_decimal(wallet.get("remaining")) >= _credit_decimal(amount): | |
| return wallet | |
| return matches[0] | |
| def _find_third_party_wallet_for_settle( | |
| wallets: list[dict[str, object]], | |
| group_id: str, | |
| model_id: str, | |
| ) -> dict[str, object] | None: | |
| matches = [wallet for wallet in wallets if _is_matching_third_party_wallet(wallet, group_id, model_id)] | |
| if not matches: | |
| return None | |
| for wallet in matches: | |
| if _credit_decimal(wallet.get("reserved")) > 0: | |
| return wallet | |
| return matches[0] | |
| def _third_party_wallet_public(wallet: dict[str, object]) -> dict[str, object]: | |
| return { | |
| "id": wallet.get("id"), | |
| "group_id": wallet.get("group_id"), | |
| "model_id": wallet.get("model_id"), | |
| "remaining": _credit_number(wallet.get("remaining")), | |
| "reserved": _credit_number(wallet.get("reserved")), | |
| "used": _credit_number(wallet.get("used")), | |
| "total_recharged": _credit_number(wallet.get("total_recharged")), | |
| "updated_at": wallet.get("updated_at"), | |
| } | |
| def _reset_package_periods(grant: dict[str, object], now: datetime | None = None) -> dict[str, object]: | |
| date_key, week_key, month_key = _period_keys(now) | |
| next_grant = dict(grant) | |
| if str(next_grant.get("daily_key") or "") != date_key: | |
| next_grant["daily_key"] = date_key | |
| next_grant["used_daily"] = 0 | |
| next_grant["reserved_daily"] = 0 | |
| if str(next_grant.get("weekly_key") or "") != week_key: | |
| next_grant["weekly_key"] = week_key | |
| next_grant["used_weekly"] = 0 | |
| next_grant["reserved_weekly"] = 0 | |
| if str(next_grant.get("monthly_key") or "") != month_key: | |
| next_grant["monthly_key"] = month_key | |
| next_grant["used_monthly"] = 0 | |
| next_grant["reserved_monthly"] = 0 | |
| return next_grant | |
| def _package_active(grant: dict[str, object]) -> bool: | |
| starts_at = str(grant.get("starts_at") or "").strip() | |
| parsed_starts_at = _parse_iso_datetime(starts_at) if starts_at else None | |
| return (parsed_starts_at is None or parsed_starts_at <= datetime.now(timezone.utc)) and _is_future_datetime(grant.get("expires_at")) | |
| def _package_available(grant: dict[str, object]) -> int | float: | |
| if not _package_active(grant): | |
| return 0 | |
| normalized = _reset_package_periods(grant) | |
| candidates: list[Decimal] = [] | |
| limit_pairs = [ | |
| ("total_limit", "used_total", "reserved_total"), | |
| ("daily_limit", "used_daily", "reserved_daily"), | |
| ("weekly_limit", "used_weekly", "reserved_weekly"), | |
| ("monthly_limit", "used_monthly", "reserved_monthly"), | |
| ] | |
| for limit_key, used_key, reserved_key in limit_pairs: | |
| limit = _to_nonnegative_int(normalized.get(limit_key)) | |
| if limit <= 0: | |
| continue | |
| candidates.append(max(Decimal("0"), Decimal(limit) - _credit_decimal(normalized.get(used_key)) - _credit_decimal(normalized.get(reserved_key)))) | |
| if not candidates: | |
| return PACKAGE_UNLIMITED_SENTINEL | |
| return _credit_number(min(candidates)) | |
| def _active_package_grants(item: dict[str, object]) -> list[dict[str, object]]: | |
| grants: list[dict[str, object]] = [] | |
| for grant in _normalize_package_grants(item.get("package_grants")): | |
| grant = _reset_package_periods(grant) | |
| if _package_active(grant): | |
| grants.append(grant) | |
| return grants | |
| def _settleable_package_grants(item: dict[str, object]) -> list[dict[str, object]]: | |
| return [_reset_package_periods(grant) for grant in _normalize_package_grants(item.get("package_grants"))] | |
| def _package_credit_available(item: dict[str, object]) -> int | float: | |
| total = Decimal("0") | |
| for grant in _active_package_grants(item): | |
| total += _credit_decimal(_package_available(grant)) | |
| return _credit_number(min(total, Decimal(PACKAGE_UNLIMITED_SENTINEL))) | |
| class AuthService: | |
| # last_used_at 的写盘节流窗口:同一密钥 60 秒内只回写一次(见 authenticate)。 | |
| _TOUCH_INTERVAL_SECONDS = 60.0 | |
| # 后台把「待写入的 last_used_at」批量落库的间隔。 | |
| _TOUCH_FLUSH_INTERVAL_SECONDS = 30.0 | |
| # 缓存未命中时回读整表的节流:未知 token 不能把事件循环拖进数据库。 | |
| _RELOAD_MIN_INTERVAL_SECONDS = 2.0 | |
| def __init__(self, storage: StorageBackend): | |
| self.storage = storage | |
| self._lock = RLock() | |
| self._items = self._load() | |
| # key_id -> 上次写 last_used_at 的时间戳;只在本进程内存里,重启即清空(无害)。 | |
| self._last_touch_ts: dict[str, float] = {} | |
| # key_id -> 待落库的 last_used_at(ISO 字符串)。 | |
| # ⚠️⚠️ 为什么不在请求里直接写库:authenticate 大量跑在**事件循环线程**上 | |
| # (async 路由里同步调 require_identity)。以前每次 touch 都要「读整表 → | |
| # 改一条 → 写整表」,图片任务线程又在并发结算积分抢同一把存储锁, | |
| # 事件循环排到锁后面 ⇒ 15 秒无响应、整站打不开(2026-10-19 事故)。 | |
| # 现在只改内存 + 记脏,由后台线程批量合并写一次。 | |
| self._pending_touch: dict[str, str] = {} | |
| self._last_reload_ts = 0.0 | |
| self._reload_requested = False | |
| self._flusher_started = False | |
| def _ensure_touch_flusher(self) -> None: | |
| """惰性启动后台 flush 线程(不在 import 期起,避免测试里凭空多线程)。""" | |
| if self._flusher_started: | |
| return | |
| self._flusher_started = True | |
| def loop() -> None: | |
| while True: | |
| time.sleep(self._TOUCH_FLUSH_INTERVAL_SECONDS) | |
| try: | |
| self._flush_reload_if_requested() | |
| self.flush_pending_touches() | |
| except Exception as exc: # pragma: no cover - 只记录,不让线程死掉 | |
| print(f"[auth-service] background flush failed: {exc}") | |
| thread = Thread(target=loop, name="auth-touch-flusher", daemon=True) | |
| thread.start() | |
| def _flush_reload_if_requested(self) -> None: | |
| """后台兜底:有未命中请求时才回读一次存储(绝不在事件循环里做)。""" | |
| with self._lock: | |
| if not self._reload_requested: | |
| return | |
| now_ts = time.time() | |
| if now_ts - self._last_reload_ts < self._RELOAD_MIN_INTERVAL_SECONDS: | |
| return | |
| self._last_reload_ts = now_ts | |
| self._reload_requested = False | |
| fallback = self._items | |
| # 读库放在锁外,避免长时间持锁。 | |
| items = self._load(fallback=fallback) | |
| with self._lock: | |
| self._items = items | |
| def flush_pending_touches(self) -> int: | |
| """把内存里攒下的 last_used_at 一次性写回存储,返回写入条数。""" | |
| with self._lock: | |
| pending = dict(self._pending_touch) | |
| self._pending_touch.clear() | |
| if not pending: | |
| return 0 | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], int]: | |
| written = 0 | |
| for index, item in enumerate(items): | |
| key_id = self._clean(item.get("id")) | |
| ts = pending.get(key_id) | |
| if not ts or item.get("last_used_at") == ts: | |
| continue | |
| next_item = dict(item) | |
| next_item["last_used_at"] = ts | |
| items[index] = next_item | |
| written += 1 | |
| return items, written | |
| try: | |
| return int(self._update_items_locked(mutate)) | |
| except Exception as exc: | |
| # 写失败就把脏数据放回去,下一轮再试(不要让统计字段拖垮请求路径)。 | |
| with self._lock: | |
| for key_id, ts in pending.items(): | |
| self._pending_touch.setdefault(key_id, ts) | |
| print(f"[auth-service] flush last_used_at failed: {exc}") | |
| return 0 | |
| def _clean(value: object) -> str: | |
| return str(value or "").strip() | |
| def _default_name(role: object) -> str: | |
| normalized = str(role or "").strip().lower() | |
| if normalized == "admin": | |
| return "管理员密钥" | |
| if normalized == "normal": | |
| return "普通用户" | |
| return "内部用户" | |
| def _normalize_item(self, raw: object) -> dict[str, object] | None: | |
| if not isinstance(raw, dict): | |
| return None | |
| role = self._clean(raw.get("role")).lower() | |
| if role not in {"admin", "user", "normal"}: | |
| return None | |
| key_hash = self._clean(raw.get("key_hash")) | |
| if not key_hash: | |
| return None | |
| item_id = self._clean(raw.get("id")) or uuid.uuid4().hex[:12] | |
| name = self._clean(raw.get("name")) | |
| if not name and not self._clean(raw.get("wechat_openid")): | |
| # 微信注册用户允许昵称为空(由用户自己在弹窗里填写),不写任何内置默认名。 | |
| name = self._default_name(role) | |
| created_at = self._clean(raw.get("created_at")) or _now_iso() | |
| last_used_at = self._clean(raw.get("last_used_at")) or None | |
| return { | |
| "id": item_id, | |
| "name": name, | |
| "role": role, | |
| "key_hash": key_hash, | |
| "enabled": bool(raw.get("enabled", True)), | |
| "account_pool_enabled": bool(raw.get("account_pool_enabled", False)), | |
| "image_memory_button_enabled": bool(raw.get("image_memory_button_enabled", False)), | |
| "third_party_image_enabled": bool(raw.get("third_party_image_enabled", False)), | |
| "third_party_image_group_ids": _normalize_string_list(raw.get("third_party_image_group_ids")), | |
| "third_party_image_model_ids": _normalize_string_list(raw.get("third_party_image_model_ids")), | |
| "third_party_image_wallets": _normalize_third_party_image_wallets(raw.get("third_party_image_wallets")), | |
| "image_model_daily_limits": _normalize_image_model_daily_limits(raw.get("image_model_daily_limits")), | |
| "image_model_daily_default_usages": _normalize_image_model_daily_limits(raw.get("image_model_daily_default_usages")), | |
| "created_at": created_at, | |
| "last_used_at": last_used_at, | |
| "credits": _credit_number(raw.get("credits")), | |
| # 签到专用账本:签到发放、且尚未被消耗掉的积分余额(见 checkin_admin_summary 注释)。 | |
| # 缺字段一律视为 0 —— 该功能上线前已发放的签到积分无法追溯,故不回溯估算。 | |
| "checkin_credit_remaining": _credit_number(raw.get("checkin_credit_remaining")), | |
| "reserved_credits": _credit_number(raw.get("reserved_credits")), | |
| "used_credits": _credit_number(raw.get("used_credits")), | |
| "total_recharged": _credit_number(raw.get("total_recharged")), | |
| "credit_grants": _normalize_credit_grants(raw.get("credit_grants")), | |
| "package_grants": _normalize_package_grants(raw.get("package_grants")), | |
| "daily_credits_date": self._clean(raw.get("daily_credits_date")), | |
| "daily_credits_used": _credit_number(raw.get("daily_credits_used")), | |
| "daily_bonus_credits": _credit_number(raw.get("daily_bonus_credits")), | |
| "daily_bonus_model_ids": _normalize_string_list(raw.get("daily_bonus_model_ids")), | |
| "daily_reserved_credits": _credit_number(raw.get("daily_reserved_credits")), | |
| "api_key_display": self._clean(raw.get("api_key_display")), | |
| "wechat_openid": self._clean(raw.get("wechat_openid")), | |
| "wechat_nickname": self._clean(raw.get("wechat_nickname")), | |
| "wechat_avatar": self._clean(raw.get("wechat_avatar")), | |
| "avatar_url": self._clean(raw.get("avatar_url")), | |
| "avatar_update_date": self._clean(raw.get("avatar_update_date")), | |
| "avatar_update_count": _to_nonnegative_int(raw.get("avatar_update_count")), | |
| "wechat_bound_at": self._clean(raw.get("wechat_bound_at")) or None, | |
| "invite_code": self._clean(raw.get("invite_code")), | |
| "invited_by": self._clean(raw.get("invited_by")), | |
| "invite_records": [dict(entry) for entry in (raw.get("invite_records") if isinstance(raw.get("invite_records"), list) else []) if isinstance(entry, dict)], | |
| "total_checkin_rewards": _to_nonnegative_int(raw.get("total_checkin_rewards")), | |
| "checkin_history": [ | |
| { | |
| "date": self._clean(entry.get("date")), | |
| "reward": _to_nonnegative_int(entry.get("reward")), | |
| "mode": "random" if self._clean(entry.get("mode")) == "random" else "checkin", | |
| "created_at": self._clean(entry.get("created_at")), | |
| } | |
| for entry in (raw.get("checkin_history") if isinstance(raw.get("checkin_history"), list) else []) | |
| if isinstance(entry, dict) and self._clean(entry.get("date")) | |
| ][-400:], | |
| } | |
| def _load(self, *, fallback: list[dict[str, object]] | None = None) -> list[dict[str, object]]: | |
| try: | |
| items = self.storage.load_auth_keys() | |
| except Exception as exc: | |
| print(f"[auth-service] load auth keys failed: {exc}") | |
| return [dict(item) for item in (fallback or [])] | |
| return self._normalize_items(items, fallback=fallback) | |
| def _normalize_items( | |
| self, | |
| items: object, | |
| *, | |
| fallback: list[dict[str, object]] | None = None, | |
| ) -> list[dict[str, object]]: | |
| if not isinstance(items, list): | |
| return [dict(item) for item in (fallback or [])] | |
| return [normalized for item in items if (normalized := self._normalize_item(item)) is not None] | |
| def _save( | |
| self, | |
| changed_items: list[dict[str, object]] | None = None, | |
| deleted_ids: list[str] | None = None, | |
| ) -> None: | |
| merge_auth_keys = getattr(self.storage, "merge_auth_keys", None) | |
| if callable(merge_auth_keys) and (changed_items is not None or deleted_ids): | |
| merge_auth_keys( | |
| [dict(item) for item in (changed_items or [])], | |
| [item_id for item_id in (self._clean(item) for item in (deleted_ids or [])) if item_id], | |
| ) | |
| return | |
| self.storage.save_auth_keys(self._items) | |
| def _reload_locked(self) -> None: | |
| self._items = self._load(fallback=self._items) | |
| def _update_items_locked( | |
| self, | |
| mutator: Callable[[list[dict[str, object]]], tuple[list[dict[str, object]], ResultT]], | |
| ) -> ResultT: | |
| update_auth_keys_locked = getattr(self.storage, "update_auth_keys_locked", None) | |
| if callable(update_auth_keys_locked): | |
| def wrapped(raw_items: list[dict[str, object]]) -> tuple[list[dict[str, object]], ResultT]: | |
| # ⚠️ 这里已经持有**存储**锁(update_auth_keys_locked 内部),再去拿鉴权锁时 | |
| # 必须保证没有「持鉴权锁等存储锁」的线程 —— 这就是 ABBA 死锁的另一半。 | |
| # 现在 authenticate 已改为「先放鉴权锁、再写存储」,加上 list_keys 不再回读 | |
| # 存储,环被打断;锁顺序统一为 存储 → 鉴权。 | |
| with self._lock: | |
| current = self._normalize_items(raw_items, fallback=self._items) | |
| self._items = current | |
| next_items, result = mutator(current) | |
| with self._lock: | |
| normalized_next = self._normalize_items(next_items, fallback=current) | |
| self._items = normalized_next | |
| return normalized_next, result | |
| return update_auth_keys_locked(wrapped) | |
| with self._lock: | |
| self._reload_locked() | |
| next_items, result = mutator(self._items) | |
| self._items = self._normalize_items(next_items, fallback=self._items) | |
| self._save() | |
| return result | |
| def _daily_state(item: dict[str, object]) -> tuple[int | float, int | float]: | |
| if item.get("role") != "normal": | |
| return 0, 0 | |
| reserved = _credit_number(item.get("daily_reserved_credits")) | |
| if AuthService._clean(item.get("daily_credits_date")) != _today_key() and reserved <= 0: | |
| return 0, 0 | |
| return _credit_number(item.get("daily_credits_used")), reserved | |
| def _daily_bonus_credits(item: dict[str, object]) -> int | float: | |
| if item.get("role") != "normal": | |
| return 0 | |
| bonus = _credit_number(item.get("daily_bonus_credits")) | |
| if bonus <= 0: | |
| return 0 | |
| if AuthService._clean(item.get("daily_credits_date")) == _today_key(): | |
| return bonus | |
| if _credit_decimal(item.get("daily_reserved_credits")) > 0: | |
| return bonus | |
| return 0 | |
| def _daily_bonus_model_ids(item: dict[str, object]) -> list[str]: | |
| if AuthService._daily_bonus_credits(item) <= 0: | |
| return [] | |
| model_ids = _normalize_string_list(item.get("daily_bonus_model_ids")) | |
| # Daily gift codes created before model scoping existed are deliberately | |
| # kept on the safest default: the built-in account pool only. | |
| return model_ids or [DAILY_GIFT_ACCOUNT_POOL_MODEL_ID] | |
| def _daily_total(item: dict[str, object]) -> int | float: | |
| if item.get("role") != "normal": | |
| return 0 | |
| return _credit_number(Decimal(config.normal_user_daily_credits) + _credit_decimal(AuthService._daily_bonus_credits(item))) | |
| def _daily_remaining(item: dict[str, object], model_id: object = "") -> int | float: | |
| used, reserved = AuthService._daily_state(item) | |
| normalized_model_id = AuthService._clean(model_id) | |
| if not normalized_model_id: | |
| return _credit_number(max(Decimal("0"), _credit_decimal(AuthService._daily_total(item)) - _credit_decimal(used) - _credit_decimal(reserved))) | |
| base_remaining = _credit_number(max(Decimal("0"), Decimal(config.normal_user_daily_credits) - _credit_decimal(used) - _credit_decimal(reserved))) | |
| if normalized_model_id not in AuthService._daily_bonus_model_ids(item): | |
| return base_remaining | |
| return _credit_number(max(Decimal("0"), _credit_decimal(AuthService._daily_total(item)) - _credit_decimal(used) - _credit_decimal(reserved))) | |
| def _reset_daily_if_needed(item: dict[str, object]) -> dict[str, object]: | |
| if item.get("role") != "normal": | |
| return item | |
| today = _today_key() | |
| if AuthService._clean(item.get("daily_credits_date")) == today: | |
| return item | |
| if _credit_decimal(item.get("daily_reserved_credits")) > 0: | |
| return item | |
| next_item = dict(item) | |
| next_item["daily_credits_date"] = today | |
| next_item["daily_credits_used"] = 0 | |
| next_item["daily_bonus_credits"] = 0 | |
| next_item["daily_bonus_model_ids"] = [] | |
| next_item["daily_reserved_credits"] = 0 | |
| return next_item | |
| def _public_item(item: dict[str, object]) -> dict[str, object]: | |
| permanent_credits = _credit_number(item.get("credits")) | |
| expiring_credits = _expiring_credit_remaining(item) | |
| paid_credits = _credit_number(_credit_decimal(permanent_credits) + _credit_decimal(expiring_credits)) | |
| package_credits = _package_credit_available(item) | |
| daily_used, daily_reserved = AuthService._daily_state(item) | |
| daily_total = AuthService._daily_total(item) | |
| daily_bonus = AuthService._daily_bonus_credits(item) | |
| daily_remaining = AuthService._daily_remaining(item) | |
| total_available = _credit_number(_credit_decimal(paid_credits) + _credit_decimal(package_credits) + _credit_decimal(daily_remaining)) if item.get("role") == "normal" else paid_credits | |
| checkin_credits, checkin_credits_expires_at = _checkin_credit_state(item) | |
| return { | |
| "id": item.get("id"), | |
| "name": item.get("name"), | |
| "role": item.get("role"), | |
| "enabled": bool(item.get("enabled", True)), | |
| "account_pool_enabled": bool(item.get("account_pool_enabled", False)), | |
| "image_memory_button_enabled": bool(item.get("image_memory_button_enabled", False)), | |
| "third_party_image_enabled": bool(item.get("third_party_image_enabled", False)), | |
| "third_party_image_group_ids": _normalize_string_list(item.get("third_party_image_group_ids")), | |
| "third_party_image_model_ids": _normalize_string_list(item.get("third_party_image_model_ids")), | |
| "third_party_image_wallets": [ | |
| _third_party_wallet_public(wallet) | |
| for wallet in _normalize_third_party_image_wallets(item.get("third_party_image_wallets")) | |
| if _to_nonnegative_int(wallet.get("remaining")) + _to_nonnegative_int(wallet.get("reserved")) + _to_nonnegative_int(wallet.get("used")) > 0 | |
| ], | |
| "image_model_daily_limits": [ | |
| _image_model_daily_limit_public(limit) | |
| for limit in _normalize_image_model_daily_limits(item.get("image_model_daily_limits")) | |
| ] if item.get("role") == "user" else [], | |
| "image_model_daily_default_usages": [ | |
| _image_model_daily_limit_public(limit) | |
| for limit in _normalize_image_model_daily_limits(item.get("image_model_daily_default_usages")) | |
| ] if item.get("role") == "user" else [], | |
| "created_at": item.get("created_at"), | |
| "last_used_at": item.get("last_used_at"), | |
| "credits": total_available, | |
| "paid_credits": paid_credits, | |
| "permanent_credits": permanent_credits, | |
| "expiring_credits": expiring_credits, | |
| "expiring_reserved_credits": _expiring_credit_reserved(item), | |
| "package_credits": package_credits if item.get("role") == "normal" else 0, | |
| "package_grants": [ | |
| { | |
| "id": grant.get("id"), | |
| "name": grant.get("name"), | |
| "remaining": _package_available(grant), | |
| "created_at": grant.get("created_at"), | |
| "starts_at": grant.get("starts_at"), | |
| "expires_at": grant.get("expires_at"), | |
| "source": grant.get("source"), | |
| "code_id": grant.get("code_id"), | |
| "daily_limit": _to_nonnegative_int(grant.get("daily_limit")), | |
| "weekly_limit": _to_nonnegative_int(grant.get("weekly_limit")), | |
| "monthly_limit": _to_nonnegative_int(grant.get("monthly_limit")), | |
| "total_limit": _to_nonnegative_int(grant.get("total_limit")), | |
| "used_total": _credit_number(grant.get("used_total")), | |
| "reserved_total": _credit_number(grant.get("reserved_total")), | |
| "used_daily": _credit_number(grant.get("used_daily")), | |
| "reserved_daily": _credit_number(grant.get("reserved_daily")), | |
| "used_weekly": _credit_number(grant.get("used_weekly")), | |
| "reserved_weekly": _credit_number(grant.get("reserved_weekly")), | |
| "used_monthly": _credit_number(grant.get("used_monthly")), | |
| "reserved_monthly": _credit_number(grant.get("reserved_monthly")), | |
| } | |
| for grant in _active_package_grants(item) | |
| if _package_available(grant) > 0 | |
| ], | |
| "credit_grants": [ | |
| { | |
| "id": grant.get("id"), | |
| "remaining": _credit_number(grant.get("remaining")), | |
| "reserved": _credit_number(grant.get("reserved")), | |
| "expires_at": grant.get("expires_at"), | |
| "source": grant.get("source"), | |
| } | |
| for grant in _active_credit_grants(item) | |
| if _credit_decimal(grant.get("remaining")) + _credit_decimal(grant.get("reserved")) > 0 | |
| ], | |
| "reserved_credits": _credit_number(item.get("reserved_credits")), | |
| "used_credits": _credit_number(item.get("used_credits")), | |
| "total_recharged": _credit_number(item.get("total_recharged")), | |
| "checkin_credit_remaining": _credit_number(item.get("checkin_credit_remaining")), | |
| # 「未使用的签到积分」= 永久形态账本 + 未过期的签到券;expires_at 为最近到期时间(空 = 永久/无)。 | |
| "checkin_credits": _credit_number(checkin_credits), | |
| "checkin_credits_expires_at": checkin_credits_expires_at, | |
| # 逐笔到期明细(按到期升序;永久形态为空列表),供前端逐行展示「x 积分 N 天后到期」。 | |
| "checkin_credit_batches": _checkin_credit_batches(item), | |
| "daily_credits": daily_total if item.get("role") == "normal" else 0, | |
| "daily_credits_remaining": daily_remaining, | |
| "daily_credits_used": daily_used, | |
| "daily_bonus_credits": daily_bonus, | |
| "daily_bonus_model_ids": AuthService._daily_bonus_model_ids(item), | |
| "daily_reserved_credits": daily_reserved, | |
| "daily_credits_date": _today_key() if item.get("role") == "normal" else "", | |
| "has_api_key_display": bool(AuthService._clean(item.get("api_key_display"))), | |
| "wechat_openid": AuthService._clean(item.get("wechat_openid")), | |
| "wechat_nickname": AuthService._clean(item.get("wechat_nickname")), | |
| "wechat_avatar": AuthService._clean(item.get("wechat_avatar")), | |
| "avatar_url": AuthService._clean(item.get("avatar_url")), | |
| "avatar_update_date": AuthService._clean(item.get("avatar_update_date")), | |
| "avatar_update_count": _to_nonnegative_int(item.get("avatar_update_count")), | |
| "wechat_bound_at": item.get("wechat_bound_at"), | |
| "invite_code": AuthService._clean(item.get("invite_code")), | |
| "invite_records": [dict(entry) for entry in item.get("invite_records", []) if isinstance(entry, dict)], | |
| "total_checkin_rewards": _to_nonnegative_int(item.get("total_checkin_rewards")), | |
| } | |
| def list_keys(self, role: AuthRole | None = None) -> list[dict[str, object]]: | |
| # ⚠️⚠️ 这里**绝对不能**再调 _reload_locked()(那会去拿存储锁)。 | |
| # 调用方常常是「已经持有存储事务锁」的商店逻辑(兑换卡密 / 每日礼包在算可用模型), | |
| # 一旦这里回头去要存储锁,就和 authenticate(持鉴权锁去写存储)形成 ABBA 死锁, | |
| # 而 authenticate 是跑在**事件循环线程**上的 ⇒ 整站彻底无响应、连接被关闭。 | |
| # 2026-10-19 事故根因。单进程内所有写入都会同步更新 self._items,缓存即权威。 | |
| with self._lock: | |
| items = [item for item in self._items if role is None or item.get("role") == role] | |
| return [self._public_item(item) for item in items] | |
| def checkin_summary(self, user_id: str) -> dict[str, object]: | |
| normalized_id = self._clean(user_id) | |
| today = _today_key() | |
| with self._lock: | |
| self._reload_locked() | |
| current = next((item for item in self._items if item.get("id") == normalized_id and item.get("role") == "normal"), None) | |
| if current is None: | |
| raise ValueError("普通用户不存在") | |
| if not self._clean(current.get("invite_code")): | |
| current["invite_code"] = _invite_code() | |
| while any(self._clean(item.get("invite_code")) == current["invite_code"] and item.get("id") != current.get("id") for item in self._items): | |
| current["invite_code"] = _invite_code() | |
| self._save(changed_items=[current]) | |
| history = [entry for entry in current.get("checkin_history", []) if isinstance(entry, dict)] | |
| own = [dict(entry) for entry in history if str(entry.get("date") or "")] | |
| records = [{"date": str(entry.get("date")), "reward": _to_nonnegative_int(entry.get("reward")), "mode": str(entry.get("mode") or "checkin")} for entry in own] | |
| users: list[dict[str, object]] = [] | |
| for item in self._items: | |
| if item.get("role") != "normal": | |
| continue | |
| entries = [entry for entry in item.get("checkin_history", []) if isinstance(entry, dict)] | |
| users.append({ | |
| "id": item.get("id"), | |
| "name": item.get("name") or "普通用户", | |
| "avatar_url": self._clean(item.get("avatar_url")) or self._clean(item.get("wechat_avatar")), | |
| "wechat_avatar": self._clean(item.get("wechat_avatar")), | |
| "history": entries, | |
| }) | |
| daily: list[dict[str, object]] = [] | |
| total: list[dict[str, object]] = [] | |
| for user in users: | |
| entries = user["history"] if isinstance(user.get("history"), list) else [] | |
| today_entries = [entry for entry in entries if isinstance(entry, dict) and str(entry.get("date") or "") == today] | |
| today_reward = sum(_to_nonnegative_int(entry.get("reward")) for entry in today_entries) | |
| today_created_at = min((str(entry.get("created_at") or "") for entry in today_entries), default="") | |
| total_reward = sum(_to_nonnegative_int(entry.get("reward")) for entry in entries if isinstance(entry, dict)) | |
| days = len({str(entry.get("date")) for entry in entries if isinstance(entry, dict) and str(entry.get("date") or "")}) | |
| if today_reward: | |
| daily.append({ | |
| "name": user["name"], | |
| "avatar_url": user.get("avatar_url") or "", | |
| "wechat_avatar": user.get("wechat_avatar") or "", | |
| "reward": today_reward, | |
| "days": days, | |
| "created_at": today_created_at, | |
| }) | |
| if total_reward: | |
| total.append({ | |
| "name": user["name"], | |
| "avatar_url": user.get("avatar_url") or "", | |
| "wechat_avatar": user.get("wechat_avatar") or "", | |
| "reward": total_reward, | |
| "days": days, | |
| }) | |
| daily.sort(key=lambda item: (str(item.get("created_at") or ""), str(item["name"]))) | |
| total.sort(key=lambda item: (-int(item["days"]), str(item["name"]))) | |
| invite_records = [] | |
| for entry in current.get("invite_records", []): | |
| if not isinstance(entry, dict): | |
| continue | |
| record = dict(entry) | |
| invitee_id = self._clean(record.get("invitee_id")) | |
| invitee = next((user for user in users if self._clean(user.get("id")) == invitee_id), None) | |
| if invitee is not None: | |
| record["invitee_name"] = invitee.get("name") or record.get("invitee_name") | |
| invite_records.append(record) | |
| total_rewards = sum(_to_nonnegative_int(entry.get("reward")) for entry in own) | |
| invite_rewards = sum(_to_nonnegative_int(entry.get("reward")) for entry in invite_records) | |
| return { | |
| "today": today, | |
| "checked_in": any(str(entry.get("date")) == today for entry in own), | |
| "records": records, | |
| "daily": daily[:100], | |
| "total": total[:100], | |
| "total_rewards": total_rewards, | |
| "invite": { | |
| "code": self._clean(current.get("invite_code")), | |
| "successful": len(invite_records), | |
| "rewards": invite_rewards, | |
| "records": invite_records, | |
| }, | |
| } | |
| def checkin_admin_summary(self) -> dict[str, object]: | |
| today = _today_key() | |
| with self._lock: | |
| self._reload_locked() | |
| users = [item for item in self._items if item.get("role") == "normal"] | |
| today_count = 0 | |
| total_rewards = 0 | |
| unused_credits = Decimal("0.00") | |
| for item in users: | |
| entries = [entry for entry in item.get("checkin_history", []) if isinstance(entry, dict)] | |
| if any(str(entry.get("date") or "") == today for entry in entries): | |
| today_count += 1 | |
| issued = sum(_to_nonnegative_int(entry.get("reward")) for entry in entries) | |
| total_rewards += issued | |
| # 「未使用」= 签到专用账本(永久形态)+ 未过期的签到券(限时形态,见 _checkin_credit_state)。 | |
| # 账本只统计本功能上线后的签到发放 —— 上线前已发放的签到积分无法追溯,不回溯估算; | |
| # 因此上线初期这个数会低于 total_rewards,属预期。 | |
| checkin_unused, _ = _checkin_credit_state(item) | |
| unused_credits += checkin_unused | |
| return { | |
| "normal_users": len(users), | |
| "today_checked_in": today_count, | |
| "total_rewards": total_rewards, | |
| "unused_credits": _credit_number(unused_credits), | |
| } | |
| def perform_checkin(self, user_id: str, *, mode: str = "checkin", reward: int = 1) -> dict[str, object]: | |
| normalized_id = self._clean(user_id) | |
| today = _today_key() | |
| normalized_mode = "random" if mode == "random" else "checkin" | |
| normalized_reward = max(1, min(1000, int(reward or 1))) | |
| # 后台「签到积分可用天数」:0 = 永久(进基础积分池);>0 则发放成带 expires_at 的限时券。 | |
| valid_days = max(0, min(3650, int(config.checkin_credit_valid_days or 0))) | |
| expires_at = (datetime.now(timezone.utc) + timedelta(days=valid_days)).isoformat() if valid_days > 0 else "" | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| history = [dict(entry) for entry in item.get("checkin_history", []) if isinstance(entry, dict)] | |
| if any(str(entry.get("date") or "") == today for entry in history): | |
| raise ValueError("今天已经签到过了") | |
| entry = {"date": today, "reward": normalized_reward, "mode": normalized_mode, "created_at": _now_iso()} | |
| if expires_at: | |
| entry["expires_at"] = expires_at | |
| next_item = dict(item) | |
| next_item["checkin_history"] = [*history[-399:], entry] | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, {"date": today, "reward": normalized_reward, "mode": normalized_mode, "expires_at": expires_at} | |
| raise ValueError("普通用户不存在") | |
| result = self._update_items_locked(mutate) | |
| self.adjust_credits(normalized_id, normalized_reward, source="checkin", expires_at=expires_at or None) | |
| return result | |
| def reconcile_checkin_credit_expiry(self, valid_days: int) -> dict[str, object]: | |
| """按后台「签到积分可用天数」重算所有普通用户**已发放**的签到积分形态。 | |
| 由后台保存该配置时调用(且仅在值真的变化时),让新设置立即对历史积分生效, | |
| 不必等用户下一次签到。语义: | |
| - valid_days > 0:把永久账本余额(`checkin_credit_remaining`)从基础积分池扣出, | |
| 转成 `expires_at = now + N 天` 的限时券(这部分原来没有到期时间,给足 N 天宽限); | |
| 已有的签到限时券**保留各自原有节奏**,到期时间 = `created_at + N 天`(缺 created_at 才退回 now + N 天)。 | |
| 若重算后已经过了到期时刻(把 N 调小到早于某些批次的发放日),这批直接作废。 | |
| - valid_days == 0:把所有签到限时券的剩余量转回基础积分池 + 永久账本,券移除。 | |
| 只处理 `source="checkin"` 的积分;卡密 / 后台加积分 / 邀请 / 套餐券一律不动。 | |
| `reserved > 0`(正在预扣)的券跳过,避免打断进行中的生成任务。 | |
| """ | |
| days = max(0, min(3650, int(valid_days or 0))) | |
| now = datetime.now(timezone.utc) | |
| expires_at = (now + timedelta(days=days)).isoformat() if days > 0 else "" | |
| converted = Decimal("0") | |
| skipped = 0 | |
| expired = 0 | |
| touched = 0 | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], None]: | |
| nonlocal converted, skipped, expired, touched | |
| next_items: list[dict[str, object]] = [] | |
| for item in items: | |
| if item.get("role") != "normal": | |
| next_items.append(item) | |
| continue | |
| grants = _normalize_credit_grants(item.get("credit_grants")) | |
| checkin_grants = [grant for grant in grants if str(grant.get("source") or "").strip() == CHECKIN_CREDIT_SOURCE] | |
| other_grants = [grant for grant in grants if str(grant.get("source") or "").strip() != CHECKIN_CREDIT_SOURCE] | |
| credits = _credit_decimal(item.get("credits")) | |
| ledger = _credit_decimal(item.get("checkin_credit_remaining")) | |
| next_credits = credits | |
| next_ledger = ledger | |
| next_checkin_grants: list[dict[str, object]] = [] | |
| for grant in checkin_grants: | |
| if _credit_decimal(grant.get("reserved")) > 0: | |
| skipped += 1 | |
| next_checkin_grants.append(grant) | |
| continue | |
| remaining = _credit_decimal(grant.get("remaining")) | |
| if remaining <= 0: | |
| continue | |
| if days > 0: | |
| # 各批次保留自己的发放节奏:到期 = 发放时刻 + N 天。 | |
| issued = _parse_iso_datetime(grant.get("created_at")) | |
| target = (issued + timedelta(days=days)) if issued is not None else (now + timedelta(days=days)) | |
| if target <= now: | |
| # 按「自发放起 N 天」的口径,这批已经过期 —— 作废(不是扣费,是自然到期)。 | |
| expired += 1 | |
| continue | |
| next_checkin_grants.append({**grant, "expires_at": target.isoformat()}) | |
| else: | |
| next_credits += remaining | |
| next_ledger += remaining | |
| converted += remaining | |
| if days > 0: | |
| movable = min(ledger, credits) | |
| if movable > 0: | |
| next_credits = credits - movable | |
| next_ledger = ledger - movable | |
| next_checkin_grants.append( | |
| { | |
| "id": uuid.uuid4().hex[:12], | |
| "remaining": _credit_number(movable), | |
| "reserved": 0, | |
| "expires_at": expires_at, | |
| "created_at": _now_iso(), | |
| "source": CHECKIN_CREDIT_SOURCE, | |
| "code_id": "", | |
| } | |
| ) | |
| converted += movable | |
| next_grants = [*other_grants, *next_checkin_grants] | |
| next_grants.sort(key=lambda grant: str(grant.get("expires_at") or "")) | |
| if next_credits == credits and next_ledger == ledger and next_grants == grants: | |
| next_items.append(item) | |
| continue | |
| next_item = dict(item) | |
| next_item["credits"] = _credit_number(next_credits) | |
| next_item["checkin_credit_remaining"] = _credit_number(next_ledger) | |
| next_item["credit_grants"] = next_grants | |
| next_items.append(next_item) | |
| touched += 1 | |
| return next_items, None | |
| self._update_items_locked(mutate) | |
| result = { | |
| "valid_days": days, | |
| "users": touched, | |
| "converted": _credit_number(converted), | |
| "expired_grants": expired, | |
| "skipped_grants": skipped, | |
| } | |
| print(f"[checkin] 签到积分有效期重算:{result}") | |
| return result | |
| def get_key(self, key_id: str, *, role: AuthRole | None = None) -> dict[str, object] | None: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| return None | |
| with self._lock: | |
| self._reload_locked() | |
| for item in self._items: | |
| if item.get("id") == normalized_id and (role is None or item.get("role") == role): | |
| return self._public_item(item) | |
| return None | |
| def _has_key_hash_locked(self, key_hash: str, *, exclude_id: str = "") -> bool: | |
| for item in self._items: | |
| item_id = self._clean(item.get("id")) | |
| if exclude_id and item_id == exclude_id: | |
| continue | |
| stored_hash = self._clean(item.get("key_hash")) | |
| if stored_hash and hmac.compare_digest(stored_hash, key_hash): | |
| return True | |
| return False | |
| def _build_key_hash_locked(self, raw_key: str, *, exclude_id: str = "") -> str: | |
| candidate = self._clean(raw_key) | |
| if not candidate: | |
| raise ValueError("请输入新的专用密钥") | |
| admin_key = self._clean(config.auth_key) | |
| if admin_key and hmac.compare_digest(candidate, admin_key): | |
| raise ValueError("这个密钥和管理员密钥冲突,请换一个") | |
| key_hash = _hash_key(candidate) | |
| if self._has_key_hash_locked(key_hash, exclude_id=exclude_id): | |
| raise ValueError("这个专用密钥已经存在,请换一个") | |
| return key_hash | |
| def _has_name_locked(self, name: str, *, role: AuthRole | None = None, exclude_id: str = "") -> bool: | |
| candidate = self._clean(name) | |
| if not candidate: | |
| return False | |
| for item in self._items: | |
| item_id = self._clean(item.get("id")) | |
| if exclude_id and item_id == exclude_id: | |
| continue | |
| if role is not None and item.get("role") != role: | |
| continue | |
| if self._clean(item.get("name")) == candidate: | |
| return True | |
| return False | |
| def _build_default_name_locked(self, role: AuthRole, *, exclude_id: str = "") -> str: | |
| base_name = self._default_name(role) | |
| if not self._has_name_locked(base_name, role=role, exclude_id=exclude_id): | |
| return base_name | |
| suffix = 2 | |
| while True: | |
| candidate = f"{base_name} {suffix}" | |
| if not self._has_name_locked(candidate, role=role, exclude_id=exclude_id): | |
| return candidate | |
| suffix += 1 | |
| def _build_name_locked(self, name: str, *, role: AuthRole, exclude_id: str = "") -> str: | |
| candidate = self._clean(name) | |
| if not candidate: | |
| return self._build_default_name_locked(role, exclude_id=exclude_id) | |
| if len(candidate) > 32: | |
| raise ValueError("昵称最多 32 个字符") | |
| if self._has_name_locked(candidate, role=role, exclude_id=exclude_id): | |
| raise ValueError("这个名称已经在使用中了,请换一个") | |
| return candidate | |
| def _build_unique_name_locked(self, name: str, *, role: AuthRole, exclude_id: str = "") -> str: | |
| candidate = self._clean(name) | |
| if not candidate: | |
| return self._build_default_name_locked(role, exclude_id=exclude_id) | |
| if not self._has_name_locked(candidate, role=role, exclude_id=exclude_id): | |
| return candidate | |
| suffix = 2 | |
| while True: | |
| next_candidate = f"{candidate} {suffix}" | |
| if not self._has_name_locked(next_candidate, role=role, exclude_id=exclude_id): | |
| return next_candidate | |
| suffix += 1 | |
| def create_key(self, *, role: AuthRole, name: str = "") -> tuple[dict[str, object], str]: | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], tuple[dict[str, object], str]]: | |
| normalized_name = self._build_name_locked(name, role=role) | |
| while True: | |
| raw_key = f"sk-{secrets.token_urlsafe(24)}" | |
| try: | |
| key_hash = self._build_key_hash_locked(raw_key) | |
| break | |
| except ValueError: | |
| continue | |
| item = { | |
| "id": uuid.uuid4().hex[:12], | |
| "name": normalized_name, | |
| "role": role, | |
| "key_hash": key_hash, | |
| "enabled": True, | |
| "account_pool_enabled": False, | |
| "image_memory_button_enabled": False, | |
| "third_party_image_enabled": False, | |
| "third_party_image_group_ids": [], | |
| "third_party_image_model_ids": [], | |
| "third_party_image_wallets": [], | |
| "image_model_daily_limits": [], | |
| "image_model_daily_default_usages": [], | |
| "created_at": _now_iso(), | |
| "last_used_at": None, | |
| "credits": config.normal_user_default_credits if role == "normal" else 0, | |
| "reserved_credits": 0, | |
| "used_credits": 0, | |
| "total_recharged": 0, | |
| "credit_grants": [], | |
| "package_grants": [], | |
| "daily_credits_date": _today_key() if role == "normal" else "", | |
| "daily_credits_used": 0, | |
| "daily_bonus_credits": 0, | |
| "daily_reserved_credits": 0, | |
| "api_key_display": raw_key if role in {"normal", "user"} else "", | |
| "wechat_openid": "", | |
| "wechat_nickname": "", | |
| "wechat_avatar": "", | |
| "wechat_bound_at": None, | |
| "checkin_history": [], | |
| "invite_code": "", | |
| "invited_by": "", | |
| "invite_records": [], | |
| "total_checkin_rewards": 0, | |
| } | |
| if role == "normal": | |
| candidate = _invite_code() | |
| while any(self._clean(existing.get("invite_code")) == candidate for existing in items): | |
| candidate = _invite_code() | |
| item["invite_code"] = candidate | |
| return [*items, item], (self._public_item(item), raw_key) | |
| return self._update_items_locked(mutate) | |
| def create_invited_normal_user(self, *, name: str, invite_code: str) -> tuple[dict[str, object], str]: | |
| normalized_name = self._clean(name) or "普通用户" | |
| normalized_invite_code = self._clean(invite_code).upper() | |
| if not normalized_invite_code: | |
| raise ValueError("请输入邀请码") | |
| invite_reward = config.checkin_invite_reward | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], tuple[dict[str, object], str]]: | |
| inviter = next((item for item in items if item.get("role") == "normal" and self._clean(item.get("invite_code")).upper() == normalized_invite_code), None) | |
| if inviter is None: | |
| raise ValueError("邀请码无效") | |
| while True: | |
| raw_key = f"sk-{secrets.token_urlsafe(24)}" | |
| key_hash = self._build_key_hash_locked(raw_key) | |
| break | |
| candidate = _invite_code() | |
| while any(self._clean(existing.get("invite_code")) == candidate for existing in items): | |
| candidate = _invite_code() | |
| item = { | |
| "id": uuid.uuid4().hex[:12], "name": self._build_unique_name_locked(normalized_name, role="normal"), "role": "normal", | |
| "key_hash": key_hash, "enabled": True, "account_pool_enabled": False, | |
| "image_memory_button_enabled": False, | |
| "third_party_image_enabled": False, "third_party_image_group_ids": [], "third_party_image_model_ids": [], | |
| "third_party_image_wallets": [], "image_model_daily_limits": [], "image_model_daily_default_usages": [], | |
| "created_at": _now_iso(), "last_used_at": None, "credits": config.normal_user_default_credits, | |
| "reserved_credits": 0, "used_credits": 0, "total_recharged": 0, "credit_grants": [], "package_grants": [], | |
| "daily_credits_date": _today_key(), "daily_credits_used": 0, "daily_bonus_credits": 0, "daily_reserved_credits": 0, | |
| "api_key_display": raw_key, "wechat_openid": "", "wechat_nickname": "", "wechat_avatar": "", "wechat_bound_at": None, | |
| "checkin_history": [], "invite_code": candidate, "invited_by": str(inviter.get("id") or ""), "invite_records": [], "total_checkin_rewards": 0, | |
| } | |
| record = {"invitee_id": item["id"], "invitee_name": item["name"], "registered_at": _now_iso(), "reward": invite_reward} | |
| inviter_next = dict(inviter) | |
| inviter_next["invite_records"] = [*([dict(entry) for entry in inviter.get("invite_records", []) if isinstance(entry, dict)])[-99:], record] | |
| next_items = list(items) | |
| next_items[next((i for i, current in enumerate(items) if current.get("id") == inviter.get("id")), 0)] = inviter_next | |
| next_items.append(item) | |
| return next_items, (self._public_item(item), raw_key) | |
| result = self._update_items_locked(mutate) | |
| inviter_id = "" | |
| created = result[0] | |
| with self._lock: | |
| self._reload_locked() | |
| inviter_id = str(next((item.get("id") for item in self._items if self._clean(item.get("invite_code")).upper() == normalized_invite_code), "") or "") | |
| if inviter_id: | |
| self.adjust_credits(inviter_id, invite_reward, source="invite") | |
| self.adjust_credits(str(created.get("id") or ""), invite_reward, source="invite") | |
| return result | |
| def attach_invite(self, user_id: str, invite_code: str) -> None: | |
| normalized_id = self._clean(user_id) | |
| normalized_code = self._clean(invite_code).upper() | |
| if not normalized_id or not normalized_code: | |
| return | |
| inviter_id = "" | |
| already = False | |
| with self._lock: | |
| self._reload_locked() | |
| target = next((item for item in self._items if item.get("id") == normalized_id and item.get("role") == "normal"), None) | |
| inviter = next((item for item in self._items if item.get("role") == "normal" and self._clean(item.get("invite_code")).upper() == normalized_code), None) | |
| if target is None or inviter is None or target.get("id") == inviter.get("id") or self._clean(target.get("invited_by")): | |
| return | |
| inviter_id = str(inviter.get("id") or "") | |
| invite_reward = config.checkin_invite_reward | |
| record = {"invitee_id": target.get("id"), "invitee_name": target.get("name"), "registered_at": _now_iso(), "reward": invite_reward} | |
| next_items = list(self._items) | |
| target_next = dict(target); target_next["invited_by"] = inviter_id | |
| inviter_next = dict(inviter); inviter_next["invite_records"] = [*([dict(entry) for entry in inviter.get("invite_records", []) if isinstance(entry, dict)])[-99:], record] | |
| next_items[next(i for i, item in enumerate(self._items) if item.get("id") == target.get("id"))] = target_next | |
| next_items[next(i for i, item in enumerate(self._items) if item.get("id") == inviter.get("id"))] = inviter_next | |
| self._items = self._normalize_items(next_items, fallback=self._items) | |
| self._save(changed_items=[target_next, inviter_next]) | |
| self.adjust_credits(inviter_id, config.checkin_invite_reward, source="invite") | |
| self.adjust_credits(normalized_id, config.checkin_invite_reward, source="invite") | |
| def get_or_create_wechat_normal_user( | |
| self, | |
| *, | |
| openid: str, | |
| nickname: str = "", | |
| avatar: str = "", | |
| allow_create: bool = True, | |
| ) -> tuple[dict[str, object], str, bool]: | |
| normalized_openid = self._clean(openid) | |
| if not normalized_openid: | |
| raise ValueError("wechat openid is required") | |
| normalized_nickname = self._clean(nickname) | |
| normalized_avatar = self._clean(avatar) | |
| now = _now_iso() | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], tuple[dict[str, object], str, bool]]: | |
| for index, item in enumerate(items): | |
| if item.get("role") != "normal" or self._clean(item.get("wechat_openid")) != normalized_openid: | |
| continue | |
| next_item = dict(item) | |
| next_item["wechat_nickname"] = normalized_nickname | |
| next_item["wechat_avatar"] = normalized_avatar | |
| next_item["wechat_bound_at"] = self._clean(next_item.get("wechat_bound_at")) or now | |
| raw_key = self._clean(next_item.get("api_key_display")) | |
| if not raw_key: | |
| while True: | |
| raw_key = f"sk-{secrets.token_urlsafe(24)}" | |
| try: | |
| next_item["key_hash"] = self._build_key_hash_locked(raw_key, exclude_id=self._clean(next_item.get("id"))) | |
| break | |
| except ValueError: | |
| continue | |
| next_item["api_key_display"] = raw_key | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, (self._public_item(next_item), raw_key, False) | |
| if not allow_create: | |
| raise ValueError("当前暂不开放新用户注册,请使用已绑定过的账号登录。") | |
| # 新注册的微信用户不预置「微信用户 xxxx」这类内置昵称,留空由用户自己设置。 | |
| normalized_name = "" | |
| while True: | |
| raw_key = f"sk-{secrets.token_urlsafe(24)}" | |
| try: | |
| key_hash = self._build_key_hash_locked(raw_key) | |
| break | |
| except ValueError: | |
| continue | |
| item = { | |
| "id": uuid.uuid4().hex[:12], | |
| "name": normalized_name, | |
| "role": "normal", | |
| "key_hash": key_hash, | |
| "enabled": True, | |
| "account_pool_enabled": False, | |
| "image_memory_button_enabled": False, | |
| "third_party_image_enabled": False, | |
| "third_party_image_group_ids": [], | |
| "third_party_image_model_ids": [], | |
| "third_party_image_wallets": [], | |
| "created_at": now, | |
| "last_used_at": None, | |
| "credits": config.normal_user_default_credits, | |
| "reserved_credits": 0, | |
| "used_credits": 0, | |
| "total_recharged": 0, | |
| "credit_grants": [], | |
| "package_grants": [], | |
| "daily_credits_date": _today_key(), | |
| "daily_credits_used": 0, | |
| "daily_bonus_credits": 0, | |
| "daily_reserved_credits": 0, | |
| "api_key_display": raw_key, | |
| "wechat_openid": normalized_openid, | |
| "wechat_nickname": normalized_nickname, | |
| "wechat_avatar": normalized_avatar, | |
| "wechat_bound_at": now, | |
| "checkin_history": [], | |
| "invite_code": _invite_code(), | |
| "invited_by": "", | |
| "invite_records": [], | |
| "total_checkin_rewards": 0, | |
| } | |
| return [*items, item], (self._public_item(item), raw_key, True) | |
| return self._update_items_locked(mutate) | |
| def get_profile(self, key_id: str) -> dict[str, object] | None: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| return None | |
| item = self.get_key(normalized_id) | |
| if item is None: | |
| return None | |
| return { | |
| "id": item.get("id"), | |
| "name": item.get("name"), | |
| "role": item.get("role"), | |
| "created_at": item.get("created_at"), | |
| "account_pool_enabled": bool(item.get("account_pool_enabled")), | |
| "third_party_image_enabled": bool(item.get("third_party_image_enabled")), | |
| "third_party_image_group_ids": _normalize_string_list(item.get("third_party_image_group_ids")), | |
| "third_party_image_model_ids": _normalize_string_list(item.get("third_party_image_model_ids")), | |
| "third_party_image_wallets": item.get("third_party_image_wallets") if isinstance(item.get("third_party_image_wallets"), list) else [], | |
| "wechat_openid": self._clean(item.get("wechat_openid")), | |
| "wechat_avatar": self._clean(item.get("wechat_avatar")), | |
| "avatar_url": self._clean(item.get("avatar_url")), | |
| "avatar_updates_remaining": max( | |
| 0, | |
| 3 - ( | |
| _to_nonnegative_int(item.get("avatar_update_count")) | |
| if self._clean(item.get("avatar_update_date")) == _today_key() | |
| else 0 | |
| ), | |
| ), | |
| "wechat_bound_at": item.get("wechat_bound_at"), | |
| "invite_code": self._clean(item.get("invite_code")), | |
| "invite_records": [dict(entry) for entry in item.get("invite_records", []) if isinstance(entry, dict)], | |
| "total_checkin_rewards": _to_nonnegative_int(item.get("total_checkin_rewards")), | |
| } | |
| def update_self_name(self, key_id: str, name: str) -> dict[str, object]: | |
| item = self.update_key(key_id, {"name": name}, role="normal") | |
| if item is None: | |
| raise ValueError("普通用户不存在") | |
| return item | |
| def update_self_avatar(self, key_id: str, avatar_url: str) -> dict[str, object]: | |
| normalized_id = self._clean(key_id) | |
| normalized_url = self._clean(avatar_url) | |
| if not normalized_id or not normalized_url: | |
| raise ValueError("头像地址无效") | |
| today = _today_key() | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| update_count = ( | |
| _to_nonnegative_int(item.get("avatar_update_count")) | |
| if self._clean(item.get("avatar_update_date")) == today | |
| else 0 | |
| ) | |
| if update_count >= 3: | |
| raise ValueError("头像每天最多修改 3 次,请明天再试") | |
| next_item = dict(item) | |
| next_item["avatar_url"] = normalized_url | |
| next_item["avatar_update_date"] = today | |
| next_item["avatar_update_count"] = update_count + 1 | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| raise ValueError("普通用户不存在") | |
| return self._update_items_locked(mutate) | |
| def get_api_key_display(self, key_id: str) -> str: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| return "" | |
| with self._lock: | |
| self._reload_locked() | |
| for item in self._items: | |
| if item.get("id") == normalized_id and item.get("role") in {"normal", "user"}: | |
| return self._clean(item.get("api_key_display")) | |
| return "" | |
| def update_key( | |
| self, | |
| key_id: str, | |
| updates: dict[str, object], | |
| *, | |
| role: AuthRole | None = None, | |
| ) -> dict[str, object] | None: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| return None | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object] | None]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id: | |
| continue | |
| if role is not None and item.get("role") != role: | |
| return items, None | |
| next_item = dict(item) | |
| next_role = str(next_item.get("role") or "user") | |
| if next_role not in {"admin", "user", "normal"}: | |
| next_role = "user" | |
| if "name" in updates and updates.get("name") is not None: | |
| next_item["name"] = self._build_name_locked( | |
| str(updates.get("name") or ""), | |
| role=next_role, # type: ignore[arg-type] | |
| exclude_id=normalized_id, | |
| ) | |
| if "enabled" in updates and updates.get("enabled") is not None: | |
| next_item["enabled"] = bool(updates.get("enabled")) | |
| if "account_pool_enabled" in updates and updates.get("account_pool_enabled") is not None: | |
| next_item["account_pool_enabled"] = bool(updates.get("account_pool_enabled")) | |
| if "image_memory_button_enabled" in updates and updates.get("image_memory_button_enabled") is not None: | |
| next_item["image_memory_button_enabled"] = bool(updates.get("image_memory_button_enabled")) | |
| if "third_party_image_enabled" in updates and updates.get("third_party_image_enabled") is not None: | |
| next_item["third_party_image_enabled"] = bool(updates.get("third_party_image_enabled")) | |
| if "third_party_image_group_ids" in updates and updates.get("third_party_image_group_ids") is not None: | |
| next_item["third_party_image_group_ids"] = _normalize_string_list(updates.get("third_party_image_group_ids")) | |
| if "third_party_image_model_ids" in updates and updates.get("third_party_image_model_ids") is not None: | |
| next_item["third_party_image_model_ids"] = _normalize_string_list(updates.get("third_party_image_model_ids")) | |
| if "image_model_daily_limits" in updates and updates.get("image_model_daily_limits") is not None: | |
| next_item["image_model_daily_limits"] = ( | |
| _normalize_image_model_daily_limits(updates.get("image_model_daily_limits")) | |
| if next_role == "user" | |
| else [] | |
| ) | |
| if "key" in updates and updates.get("key") is not None: | |
| raw_key = str(updates.get("key") or "") | |
| next_item["key_hash"] = self._build_key_hash_locked(raw_key, exclude_id=normalized_id) | |
| if next_role in {"normal", "user"}: | |
| next_item["api_key_display"] = self._clean(raw_key) | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| return items, None | |
| return self._update_items_locked(mutate) | |
| def reserve_internal_image_model_daily_limit( | |
| self, | |
| key_id: str, | |
| *, | |
| model_id: str, | |
| amount: int, | |
| ) -> dict[str, object]: | |
| normalized_id = self._clean(key_id) | |
| normalized_model_id = self._clean(model_id) | |
| normalized_amount = max(0, int(amount or 0)) | |
| if not normalized_id or not normalized_model_id or normalized_amount <= 0: | |
| return {"_internal_image_model_daily_limit_reservation": {"amount": 0}} | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "user": | |
| continue | |
| limits = _normalize_image_model_daily_limits(item.get("image_model_daily_limits")) | |
| limit = next((entry for entry in limits if self._clean(entry.get("model_id")) == normalized_model_id), None) | |
| source_field = "image_model_daily_limits" | |
| reservation_scope = "user" | |
| limit_scope = "per_user" | |
| if limit is None: | |
| default_limit = config.get_internal_image_model_daily_limit(normalized_model_id) | |
| if default_limit is None: | |
| return items, {"_internal_image_model_daily_limit_reservation": {"amount": 0}} | |
| limit_scope = str(default_limit.get("limit_scope") or "per_user").strip().lower() | |
| if limit_scope == "total": | |
| total_result = config.reserve_internal_image_model_total_daily_limit(normalized_model_id, normalized_amount) | |
| total_reservation = total_result.get("_internal_image_model_daily_limit_reservation") if isinstance(total_result, dict) else {} | |
| return items, { | |
| "_internal_image_model_daily_limit_reservation": total_reservation | |
| if isinstance(total_reservation, dict) | |
| else {"amount": 0} | |
| } | |
| source_field = "image_model_daily_default_usages" | |
| reservation_scope = "default_user" | |
| limits = _normalize_image_model_daily_limits(item.get(source_field)) | |
| limit = next((entry for entry in limits if self._clean(entry.get("model_id")) == normalized_model_id), None) | |
| if limit is None: | |
| limit = { | |
| "model_id": normalized_model_id, | |
| "daily_limit": _to_nonnegative_int(default_limit.get("daily_limit")), | |
| "date": _today_key(), | |
| "used": 0, | |
| "reserved": 0, | |
| "updated_at": "", | |
| } | |
| limits.append(limit) | |
| else: | |
| limit["daily_limit"] = _to_nonnegative_int(default_limit.get("daily_limit")) | |
| limit = _reset_image_model_daily_limit(limit) | |
| daily_limit = _to_nonnegative_int(limit.get("daily_limit")) | |
| used = _to_nonnegative_int(limit.get("used")) | |
| reserved = _to_nonnegative_int(limit.get("reserved")) | |
| remaining = max(0, daily_limit - used - reserved) | |
| if normalized_amount > remaining: | |
| if remaining <= 0: | |
| raise ValueError(f"{normalized_model_id} 今日使用数量已达上限") | |
| raise ValueError(f"{normalized_model_id} 今日剩余 {remaining} 张,不足以生成 {normalized_amount} 张") | |
| limit["reserved"] = reserved + normalized_amount | |
| limit["date"] = str(limit.get("date") or "").strip() or _today_key() | |
| limit["updated_at"] = _now_iso() | |
| next_limits = [ | |
| limit if self._clean(entry.get("model_id")) == normalized_model_id else entry | |
| for entry in limits | |
| ] | |
| next_item = dict(item) | |
| next_item[source_field] = next_limits | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| result = self._public_item(next_item) | |
| result["_internal_image_model_daily_limit_reservation"] = { | |
| "model_id": normalized_model_id, | |
| "amount": normalized_amount, | |
| "scope": reservation_scope, | |
| "limit_scope": limit_scope, | |
| } | |
| return next_items, result | |
| return items, {"_internal_image_model_daily_limit_reservation": {"amount": 0}} | |
| return self._update_items_locked(mutate) | |
| def consume_internal_image_model_daily_limit( | |
| self, | |
| key_id: str, | |
| *, | |
| model_id: str, | |
| amount: int, | |
| scope: str = "", | |
| ) -> dict[str, object] | None: | |
| normalized_id = self._clean(key_id) | |
| normalized_model_id = self._clean(model_id) | |
| normalized_amount = max(0, int(amount or 0)) | |
| if not normalized_id or not normalized_model_id or normalized_amount <= 0: | |
| return None | |
| normalized_scope = self._clean(scope) | |
| if normalized_scope == "default_total": | |
| return config.consume_internal_image_model_total_daily_limit(normalized_model_id, normalized_amount) | |
| source_field = "image_model_daily_default_usages" if normalized_scope == "default_user" else "image_model_daily_limits" | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object] | None]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "user": | |
| continue | |
| limits = _normalize_image_model_daily_limits(item.get(source_field)) | |
| changed = False | |
| for limit in limits: | |
| if self._clean(limit.get("model_id")) != normalized_model_id: | |
| continue | |
| limit = _reset_image_model_daily_limit(limit) | |
| reserved = _to_nonnegative_int(limit.get("reserved")) | |
| consumed = min(reserved, normalized_amount) | |
| limit["reserved"] = reserved - consumed | |
| limit["used"] = _to_nonnegative_int(limit.get("used")) + consumed | |
| limit["updated_at"] = _now_iso() | |
| changed = True | |
| next_limits = [ | |
| limit if self._clean(entry.get("model_id")) == normalized_model_id else entry | |
| for entry in limits | |
| ] | |
| next_item = dict(item) | |
| next_item[source_field] = next_limits | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| if not changed: | |
| return items, None | |
| return items, None | |
| return self._update_items_locked(mutate) | |
| def refund_internal_image_model_daily_limit( | |
| self, | |
| key_id: str, | |
| *, | |
| model_id: str, | |
| amount: int, | |
| scope: str = "", | |
| ) -> dict[str, object] | None: | |
| normalized_id = self._clean(key_id) | |
| normalized_model_id = self._clean(model_id) | |
| normalized_amount = max(0, int(amount or 0)) | |
| if not normalized_id or not normalized_model_id or normalized_amount <= 0: | |
| return None | |
| normalized_scope = self._clean(scope) | |
| if normalized_scope == "default_total": | |
| return config.refund_internal_image_model_total_daily_limit(normalized_model_id, normalized_amount) | |
| source_field = "image_model_daily_default_usages" if normalized_scope == "default_user" else "image_model_daily_limits" | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object] | None]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "user": | |
| continue | |
| limits = _normalize_image_model_daily_limits(item.get(source_field)) | |
| for limit in limits: | |
| if self._clean(limit.get("model_id")) != normalized_model_id: | |
| continue | |
| limit = _reset_image_model_daily_limit(limit) | |
| reserved = _to_nonnegative_int(limit.get("reserved")) | |
| refunded = min(reserved, normalized_amount) | |
| limit["reserved"] = reserved - refunded | |
| limit["updated_at"] = _now_iso() | |
| next_limits = [ | |
| limit if self._clean(entry.get("model_id")) == normalized_model_id else entry | |
| for entry in limits | |
| ] | |
| next_item = dict(item) | |
| next_item[source_field] = next_limits | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| return items, None | |
| return self._update_items_locked(mutate) | |
| def delete_key(self, key_id: str, *, role: AuthRole | None = None) -> bool: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| return False | |
| # Database-backed stores can delete one row directly. Avoid loading, | |
| # normalizing, and rewriting every user record: a single malformed | |
| # legacy row should not prevent an unrelated user from being deleted. | |
| merge_auth_keys = getattr(self.storage, "merge_auth_keys", None) | |
| if callable(merge_auth_keys): | |
| with self._lock: | |
| items = self._load(fallback=self._items) | |
| target = next( | |
| ( | |
| item | |
| for item in items | |
| if item.get("id") == normalized_id | |
| and (role is None or item.get("role") == role) | |
| ), | |
| None, | |
| ) | |
| if target is None: | |
| self._items = items | |
| return False | |
| merge_auth_keys([], [normalized_id]) | |
| self._items = [item for item in items if item.get("id") != normalized_id] | |
| return True | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], bool]: | |
| next_items = [ | |
| item | |
| for item in items | |
| if not (item.get("id") == normalized_id and (role is None or item.get("role") == role)) | |
| ] | |
| return next_items, len(next_items) < len(items) | |
| return self._update_items_locked(mutate) | |
| def batch_mutate_keys( | |
| self, | |
| key_ids: list[str], | |
| *, | |
| action: Literal["enable", "disable", "delete"], | |
| role: AuthRole, | |
| ) -> dict[str, object]: | |
| """Enable, disable, or delete multiple users in one storage transaction.""" | |
| normalized_ids = list( | |
| dict.fromkeys( | |
| item_id | |
| for item_id in (self._clean(key_id) for key_id in key_ids) | |
| if item_id | |
| ) | |
| ) | |
| if not normalized_ids: | |
| raise ValueError("请至少选择一个用户") | |
| if action not in {"enable", "disable", "delete"}: | |
| raise ValueError("不支持的批量操作") | |
| if role not in {"user", "normal"}: | |
| raise ValueError("不支持的用户类型") | |
| selected_ids = set(normalized_ids) | |
| with self._lock: | |
| items = self._load(fallback=self._items) | |
| matched_ids = { | |
| str(item.get("id") or "") | |
| for item in items | |
| if item.get("id") in selected_ids and item.get("role") == role | |
| } | |
| if action == "delete": | |
| next_items = [item for item in items if str(item.get("id") or "") not in matched_ids] | |
| if matched_ids: | |
| self._items = next_items | |
| self._save(changed_items=[], deleted_ids=list(matched_ids)) | |
| else: | |
| self._items = items | |
| changed_count = len(matched_ids) | |
| else: | |
| enabled = action == "enable" | |
| changed_items: list[dict[str, object]] = [] | |
| next_items: list[dict[str, object]] = [] | |
| for item in items: | |
| item_id = str(item.get("id") or "") | |
| if item_id not in matched_ids or bool(item.get("enabled")) == enabled: | |
| next_items.append(item) | |
| continue | |
| next_item = dict(item) | |
| next_item["enabled"] = enabled | |
| next_items.append(next_item) | |
| changed_items.append(next_item) | |
| self._items = next_items | |
| if changed_items: | |
| self._save(changed_items=changed_items, deleted_ids=[]) | |
| changed_count = len(changed_items) | |
| return { | |
| "action": action, | |
| "requested": len(normalized_ids), | |
| "matched": len(matched_ids), | |
| "changed": changed_count, | |
| } | |
| def batch_update_image_permissions( | |
| self, | |
| updates: list[dict[str, object]], | |
| *, | |
| role: AuthRole, | |
| ) -> dict[str, object]: | |
| """Update image routing permissions for many users in one storage transaction.""" | |
| if role not in {"user", "normal"}: | |
| raise ValueError("不支持的用户类型") | |
| normalized_updates = { | |
| self._clean(update.get("id")): update | |
| for update in updates | |
| if isinstance(update, dict) and self._clean(update.get("id")) | |
| } | |
| if not normalized_updates: | |
| raise ValueError("请至少选择一个用户") | |
| with self._lock: | |
| items = self._load(fallback=self._items) | |
| changed_items: list[dict[str, object]] = [] | |
| next_items: list[dict[str, object]] = [] | |
| matched = 0 | |
| for item in items: | |
| item_id = self._clean(item.get("id")) | |
| update = normalized_updates.get(item_id) | |
| if update is None or item.get("role") != role: | |
| next_items.append(item) | |
| continue | |
| matched += 1 | |
| next_item = dict(item) | |
| if "third_party_image_enabled" in update: | |
| next_item["third_party_image_enabled"] = bool(update.get("third_party_image_enabled")) | |
| if "third_party_image_group_ids" in update: | |
| next_item["third_party_image_group_ids"] = _normalize_string_list(update.get("third_party_image_group_ids")) | |
| if "third_party_image_model_ids" in update: | |
| next_item["third_party_image_model_ids"] = _normalize_string_list(update.get("third_party_image_model_ids")) | |
| next_items.append(next_item) | |
| if next_item != item: | |
| changed_items.append(next_item) | |
| self._items = next_items | |
| if changed_items: | |
| self._save(changed_items=changed_items, deleted_ids=[]) | |
| return { | |
| "requested": len(normalized_updates), | |
| "matched": matched, | |
| "changed": len(changed_items), | |
| } | |
| def credit_summary(self, identity: dict[str, object]) -> dict[str, object]: | |
| role = self._clean(identity.get("role")).lower() | |
| if role != "normal": | |
| return { | |
| "role": role, | |
| "metered": False, | |
| "credits": None, | |
| "paid_credits": None, | |
| "reserved_credits": 0, | |
| "used_credits": None, | |
| "total_recharged": None, | |
| "permanent_credits": None, | |
| "expiring_credits": 0, | |
| "expiring_reserved_credits": 0, | |
| "credit_grants": [], | |
| "package_credits": 0, | |
| "package_grants": [], | |
| "daily_credits": 0, | |
| "daily_credits_remaining": 0, | |
| "daily_credits_used": None, | |
| "daily_bonus_credits": 0, | |
| "daily_reserved_credits": 0, | |
| "daily_credits_date": "", | |
| "checkin_credits": 0, | |
| "checkin_credits_expires_at": "", | |
| "checkin_credit_batches": [], | |
| "third_party_image_enabled": False, | |
| "third_party_image_wallets": [], | |
| } | |
| item = self.get_key(self._clean(identity.get("id")), role="normal") | |
| if item is None: | |
| return { | |
| "role": role, | |
| "metered": True, | |
| "credits": 0, | |
| "paid_credits": 0, | |
| "reserved_credits": 0, | |
| "used_credits": 0, | |
| "total_recharged": 0, | |
| "permanent_credits": 0, | |
| "expiring_credits": 0, | |
| "expiring_reserved_credits": 0, | |
| "credit_grants": [], | |
| "package_credits": 0, | |
| "package_grants": [], | |
| "daily_credits": config.normal_user_daily_credits, | |
| "daily_credits_remaining": 0, | |
| "daily_credits_used": 0, | |
| "daily_bonus_credits": 0, | |
| "daily_reserved_credits": 0, | |
| "daily_credits_date": _today_key(), | |
| "checkin_credits": 0, | |
| "checkin_credits_expires_at": "", | |
| "checkin_credit_batches": [], | |
| "third_party_image_enabled": False, | |
| "third_party_image_wallets": [], | |
| } | |
| return self.credit_summary_from_public_item(item) | |
| def credit_summary_from_public_item(self, item: dict[str, object]) -> dict[str, object]: | |
| return { | |
| "role": "normal", | |
| "metered": True, | |
| "credits": _credit_number(item.get("credits")), | |
| "paid_credits": _credit_number(item.get("paid_credits")), | |
| "reserved_credits": _credit_number(item.get("reserved_credits")), | |
| "used_credits": _credit_number(item.get("used_credits")), | |
| "total_recharged": _credit_number(item.get("total_recharged")), | |
| "permanent_credits": _credit_number(item.get("permanent_credits")), | |
| "expiring_credits": _credit_number(item.get("expiring_credits")), | |
| "expiring_reserved_credits": _credit_number(item.get("expiring_reserved_credits")), | |
| "credit_grants": item.get("credit_grants") if isinstance(item.get("credit_grants"), list) else [], | |
| "package_credits": _credit_number(item.get("package_credits")), | |
| "package_grants": item.get("package_grants") if isinstance(item.get("package_grants"), list) else [], | |
| "daily_credits": _credit_number(item.get("daily_credits")), | |
| "daily_credits_remaining": _credit_number(item.get("daily_credits_remaining")), | |
| "daily_credits_used": _credit_number(item.get("daily_credits_used")), | |
| "daily_bonus_credits": _credit_number(item.get("daily_bonus_credits")), | |
| "daily_reserved_credits": _credit_number(item.get("daily_reserved_credits")), | |
| "daily_credits_date": self._clean(item.get("daily_credits_date")), | |
| "checkin_credits": _credit_number(item.get("checkin_credits")), | |
| "checkin_credits_expires_at": self._clean(item.get("checkin_credits_expires_at")), | |
| "checkin_credit_batches": item.get("checkin_credit_batches") if isinstance(item.get("checkin_credit_batches"), list) else [], | |
| "third_party_image_enabled": bool(item.get("third_party_image_enabled")), | |
| "third_party_image_wallets": item.get("third_party_image_wallets") if isinstance(item.get("third_party_image_wallets"), list) else [], | |
| } | |
| def _adjust_credits_in_items( | |
| self, | |
| items: list[dict[str, object]], | |
| key_id: str, | |
| delta: int, | |
| *, | |
| note: str = "", | |
| expires_at: str | None = None, | |
| source: str = "", | |
| code_id: str = "", | |
| ) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| amount = int(delta) | |
| for index, item in enumerate(items): | |
| if item.get("id") != key_id or item.get("role") != "normal": | |
| continue | |
| next_item = dict(item) | |
| credits = _credit_decimal(item.get("credits")) | |
| if amount > 0 and _is_future_datetime(expires_at): | |
| grants = _normalize_credit_grants(item.get("credit_grants")) | |
| grants.append( | |
| { | |
| "id": uuid.uuid4().hex[:12], | |
| "remaining": amount, | |
| "reserved": 0, | |
| "expires_at": str(expires_at or "").strip(), | |
| "created_at": _now_iso(), | |
| "source": self._clean(source or note), | |
| "code_id": self._clean(code_id), | |
| } | |
| ) | |
| grants.sort(key=lambda grant: str(grant.get("expires_at") or "")) | |
| next_item["credit_grants"] = grants | |
| next_item["credits"] = _credit_number(credits) | |
| else: | |
| next_credits = credits + Decimal(amount) | |
| if next_credits < 0: | |
| raise ValueError("用户剩余次数不足,不能扣成负数") | |
| next_item["credits"] = _credit_number(next_credits) | |
| # 签到账本:只有落进基础积分池的变动才记账。 | |
| # 发放(source=checkin)累加;手动扣分按「签到积分最先被消耗」扣减。 | |
| checkin_remaining = _credit_decimal(item.get("checkin_credit_remaining")) | |
| if amount > 0: | |
| if self._clean(source) == "checkin": | |
| checkin_remaining += Decimal(amount) | |
| elif amount < 0: | |
| checkin_remaining = max(Decimal("0"), checkin_remaining + Decimal(amount)) | |
| next_item["checkin_credit_remaining"] = _credit_number(checkin_remaining) | |
| if amount > 0: | |
| next_item["total_recharged"] = _credit_number(_credit_decimal(item.get("total_recharged")) + Decimal(amount)) | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| raise ValueError("普通用户不存在") | |
| def adjust_credits(self, key_id: str, delta: int, *, note: str = "", expires_at: str | None = None, source: str = "", code_id: str = "") -> dict[str, object]: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| raise ValueError("用户不存在") | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| return self._adjust_credits_in_items( | |
| items, | |
| normalized_id, | |
| int(delta), | |
| note=note, | |
| expires_at=expires_at, | |
| source=source, | |
| code_id=code_id, | |
| ) | |
| return self._update_items_locked(mutate) | |
| def _add_third_party_image_wallet_in_items( | |
| self, | |
| items: list[dict[str, object]], | |
| key_id: str, | |
| *, | |
| group_id: str, | |
| model_id: str, | |
| amount: int, | |
| auto_enable: bool = True, | |
| ) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| normalized_id = self._clean(key_id) | |
| normalized_group_id = self._clean(group_id) | |
| normalized_model_id = self._clean(model_id) | |
| normalized_amount = max(1, _to_nonnegative_int(amount)) | |
| if not normalized_group_id: | |
| raise ValueError("image API group is required") | |
| if not normalized_model_id: | |
| raise ValueError("image API model is required") | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| wallets = _normalize_third_party_image_wallets(item.get("third_party_image_wallets")) | |
| wallet_key = f"{normalized_group_id}:{normalized_model_id}" | |
| wallet = next((entry for entry in wallets if self._clean(entry.get("id")) == wallet_key), None) | |
| if wallet is None: | |
| wallet = { | |
| "id": wallet_key, | |
| "group_id": normalized_group_id, | |
| "model_id": normalized_model_id, | |
| "remaining": 0, | |
| "reserved": 0, | |
| "used": 0, | |
| "total_recharged": 0, | |
| "updated_at": "", | |
| } | |
| wallets.append(wallet) | |
| wallet["remaining"] = _credit_number(_credit_decimal(wallet.get("remaining")) + _credit_decimal(normalized_amount)) | |
| wallet["total_recharged"] = _credit_number(_credit_decimal(wallet.get("total_recharged")) + _credit_decimal(normalized_amount)) | |
| wallet["updated_at"] = _now_iso() | |
| next_item = dict(item) | |
| next_item["third_party_image_wallets"] = wallets | |
| if auto_enable: | |
| next_item["third_party_image_enabled"] = True | |
| if normalized_model_id == THIRD_PARTY_IMAGE_GROUP_WALLET_MODEL_ID: | |
| group_ids = set(_normalize_string_list(next_item.get("third_party_image_group_ids"))) | |
| group_ids.add(normalized_group_id) | |
| next_item["third_party_image_group_ids"] = sorted(group_ids) | |
| next_item["third_party_image_model_ids"] = [] | |
| else: | |
| model_ids = set(_normalize_string_list(next_item.get("third_party_image_model_ids"))) | |
| model_ids.add(f"{normalized_group_id}:{normalized_model_id}") | |
| next_item["third_party_image_model_ids"] = sorted(model_ids) | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| raise ValueError("normal user not found") | |
| def adjust_third_party_image_wallet( | |
| self, | |
| key_id: str, | |
| *, | |
| group_id: str, | |
| model_id: str, | |
| delta: int, | |
| source: str = "", | |
| code_id: str = "", | |
| auto_enable: bool = True, | |
| ) -> dict[str, object]: | |
| normalized_id = self._clean(key_id) | |
| normalized_group_id = self._clean(group_id) | |
| normalized_model_id = self._clean(model_id) | |
| amount = int(delta) | |
| if not normalized_id: | |
| raise ValueError("user not found") | |
| if not normalized_group_id: | |
| raise ValueError("image API group is required") | |
| if not normalized_model_id: | |
| raise ValueError("image API model is required") | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| wallets = _normalize_third_party_image_wallets(item.get("third_party_image_wallets")) | |
| wallet_key = f"{normalized_group_id}:{normalized_model_id}" | |
| wallet = next((entry for entry in wallets if self._clean(entry.get("id")) == wallet_key), None) | |
| if wallet is None: | |
| wallet = { | |
| "id": wallet_key, | |
| "group_id": normalized_group_id, | |
| "model_id": normalized_model_id, | |
| "remaining": 0, | |
| "reserved": 0, | |
| "used": 0, | |
| "total_recharged": 0, | |
| "updated_at": "", | |
| } | |
| wallets.append(wallet) | |
| next_remaining = _credit_decimal(wallet.get("remaining")) + Decimal(amount) | |
| if next_remaining < Decimal("0"): | |
| raise ValueError("4K image balance is not enough") | |
| wallet["remaining"] = _credit_number(next_remaining) | |
| wallet["updated_at"] = _now_iso() | |
| if amount > 0: | |
| wallet["total_recharged"] = _credit_number(_credit_decimal(wallet.get("total_recharged")) + Decimal(amount)) | |
| next_item = dict(item) | |
| next_item["third_party_image_wallets"] = wallets | |
| if auto_enable: | |
| next_item["third_party_image_enabled"] = True | |
| if normalized_model_id == THIRD_PARTY_IMAGE_GROUP_WALLET_MODEL_ID: | |
| group_ids = set(_normalize_string_list(next_item.get("third_party_image_group_ids"))) | |
| group_ids.add(normalized_group_id) | |
| next_item["third_party_image_group_ids"] = sorted(group_ids) | |
| next_item["third_party_image_model_ids"] = [] | |
| else: | |
| model_ids = set(_normalize_string_list(next_item.get("third_party_image_model_ids"))) | |
| model_ids.add(f"{normalized_group_id}:{normalized_model_id}") | |
| next_item["third_party_image_model_ids"] = sorted(model_ids) | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| raise ValueError("normal user not found") | |
| return self._update_items_locked(mutate) | |
| def reserve_third_party_image_wallet( | |
| self, | |
| key_id: str, | |
| *, | |
| group_id: str, | |
| model_id: str, | |
| amount: int | float, | |
| ) -> dict[str, object]: | |
| normalized_id = self._clean(key_id) | |
| normalized_group_id = self._clean(group_id) | |
| normalized_model_id = self._clean(model_id) | |
| normalized_amount = _credit_decimal(amount) | |
| if normalized_amount <= 0: | |
| raise ValueError("image API model balance amount must be positive") | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| wallets = _normalize_third_party_image_wallets(item.get("third_party_image_wallets")) | |
| wallet = _find_third_party_wallet_for_reserve(wallets, normalized_group_id, normalized_model_id, normalized_amount) | |
| if wallet is None: | |
| raise ValueError("image API model balance is not enough") | |
| remaining = _credit_decimal(wallet.get("remaining")) | |
| if remaining < normalized_amount: | |
| raise ValueError("image API model balance is not enough") | |
| wallet["remaining"] = _credit_number(remaining - normalized_amount) | |
| wallet["reserved"] = _credit_number(_credit_decimal(wallet.get("reserved")) + normalized_amount) | |
| wallet["updated_at"] = _now_iso() | |
| next_item = dict(item) | |
| next_item["third_party_image_wallets"] = wallets | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| result = self._public_item(next_item) | |
| result["_third_party_image_reservation"] = { | |
| "group_id": normalized_group_id, | |
| "model_id": normalized_model_id, | |
| "amount": _credit_number(normalized_amount), | |
| } | |
| return next_items, result | |
| raise ValueError("normal user not found") | |
| return self._update_items_locked(mutate) | |
| def consume_third_party_image_wallet( | |
| self, | |
| key_id: str, | |
| *, | |
| group_id: str, | |
| model_id: str, | |
| amount: int | float, | |
| ) -> dict[str, object] | None: | |
| normalized_id = self._clean(key_id) | |
| normalized_group_id = self._clean(group_id) | |
| normalized_model_id = self._clean(model_id) | |
| normalized_amount = _credit_decimal(amount) | |
| if normalized_amount <= 0: | |
| return None | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object] | None]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| wallets = _normalize_third_party_image_wallets(item.get("third_party_image_wallets")) | |
| wallet = _find_third_party_wallet_for_settle(wallets, normalized_group_id, normalized_model_id) | |
| if wallet is not None: | |
| reserved = _credit_decimal(wallet.get("reserved")) | |
| consumed = min(reserved, normalized_amount) | |
| wallet["reserved"] = _credit_number(reserved - consumed) | |
| wallet["used"] = _credit_number(_credit_decimal(wallet.get("used")) + consumed) | |
| wallet["updated_at"] = _now_iso() | |
| next_item = dict(item) | |
| next_item["third_party_image_wallets"] = wallets | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| return items, None | |
| return items, None | |
| return self._update_items_locked(mutate) | |
| def refund_third_party_image_wallet( | |
| self, | |
| key_id: str, | |
| *, | |
| group_id: str, | |
| model_id: str, | |
| amount: int | float, | |
| ) -> dict[str, object] | None: | |
| normalized_id = self._clean(key_id) | |
| normalized_group_id = self._clean(group_id) | |
| normalized_model_id = self._clean(model_id) | |
| normalized_amount = _credit_decimal(amount) | |
| if normalized_amount <= 0: | |
| return None | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object] | None]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| wallets = _normalize_third_party_image_wallets(item.get("third_party_image_wallets")) | |
| wallet = _find_third_party_wallet_for_settle(wallets, normalized_group_id, normalized_model_id) | |
| if wallet is not None: | |
| reserved = _credit_decimal(wallet.get("reserved")) | |
| refunded = min(reserved, normalized_amount) | |
| wallet["reserved"] = _credit_number(reserved - refunded) | |
| wallet["remaining"] = _credit_number(_credit_decimal(wallet.get("remaining")) + refunded) | |
| wallet["updated_at"] = _now_iso() | |
| next_item = dict(item) | |
| next_item["third_party_image_wallets"] = wallets | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| return items, None | |
| return items, None | |
| return self._update_items_locked(mutate) | |
| def _add_daily_credits_in_items( | |
| self, | |
| items: list[dict[str, object]], | |
| key_id: str, | |
| amount: int, | |
| *, | |
| source: str = "", | |
| code_id: str = "", | |
| model_ids: list[str] | None = None, | |
| ) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| normalized_amount = max(1, _to_nonnegative_int(amount)) | |
| today = _today_key() | |
| for index, item in enumerate(items): | |
| if item.get("id") != key_id or item.get("role") != "normal": | |
| continue | |
| normalized_item = self._reset_daily_if_needed(dict(item)) | |
| next_item = dict(normalized_item) | |
| next_item["daily_credits_date"] = today | |
| next_item["daily_bonus_credits"] = _to_nonnegative_int(normalized_item.get("daily_bonus_credits")) + normalized_amount | |
| normalized_model_ids = _normalize_string_list(model_ids) or [DAILY_GIFT_ACCOUNT_POOL_MODEL_ID] | |
| merged_model_ids = set(self._daily_bonus_model_ids(normalized_item)) | |
| merged_model_ids.update(normalized_model_ids) | |
| next_item["daily_bonus_model_ids"] = sorted(merged_model_ids) | |
| third_party_model_ids = { | |
| model_id | |
| for model_id in normalized_model_ids | |
| if model_id != DAILY_GIFT_ACCOUNT_POOL_MODEL_ID and ":" in model_id | |
| } | |
| if third_party_model_ids: | |
| next_item["third_party_image_enabled"] = True | |
| permitted_model_ids = set(_normalize_string_list(next_item.get("third_party_image_model_ids"))) | |
| permitted_model_ids.update(third_party_model_ids) | |
| next_item["third_party_image_model_ids"] = sorted(permitted_model_ids) | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| raise ValueError("普通用户不存在") | |
| def add_daily_credits( | |
| self, | |
| key_id: str, | |
| amount: int, | |
| *, | |
| source: str = "", | |
| code_id: str = "", | |
| model_ids: list[str] | None = None, | |
| ) -> dict[str, object]: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| raise ValueError("用户不存在") | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| return self._add_daily_credits_in_items( | |
| items, | |
| normalized_id, | |
| int(amount), | |
| source=source, | |
| code_id=code_id, | |
| model_ids=model_ids, | |
| ) | |
| return self._update_items_locked(mutate) | |
| def _add_package_grant_in_items( | |
| self, | |
| items: list[dict[str, object]], | |
| key_id: str, | |
| package_plan: dict[str, object], | |
| *, | |
| source: str = "", | |
| code_id: str = "", | |
| ) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| duration_days = max(1, _to_nonnegative_int(package_plan.get("duration_days")) or 1) | |
| now = datetime.now(timezone.utc) | |
| grant = { | |
| "id": uuid.uuid4().hex[:12], | |
| "name": self._clean(package_plan.get("name")) or "套餐", | |
| "created_at": now.isoformat(), | |
| "starts_at": now.isoformat(), | |
| "expires_at": (now + timedelta(days=duration_days)).isoformat(), | |
| "source": self._clean(source), | |
| "code_id": self._clean(code_id), | |
| "daily_limit": _to_nonnegative_int(package_plan.get("daily_limit")), | |
| "weekly_limit": _to_nonnegative_int(package_plan.get("weekly_limit")), | |
| "monthly_limit": _to_nonnegative_int(package_plan.get("monthly_limit")), | |
| "total_limit": _to_nonnegative_int(package_plan.get("total_limit")), | |
| "used_total": 0, | |
| "reserved_total": 0, | |
| "used_daily": 0, | |
| "reserved_daily": 0, | |
| "used_weekly": 0, | |
| "reserved_weekly": 0, | |
| "used_monthly": 0, | |
| "reserved_monthly": 0, | |
| } | |
| if max( | |
| _to_nonnegative_int(grant.get("daily_limit")), | |
| _to_nonnegative_int(grant.get("weekly_limit")), | |
| _to_nonnegative_int(grant.get("monthly_limit")), | |
| _to_nonnegative_int(grant.get("total_limit")), | |
| ) <= 0: | |
| raise ValueError("套餐卡密至少需要设置一个日/周/月/总限额") | |
| grant = _reset_package_periods(grant) | |
| for index, item in enumerate(items): | |
| if item.get("id") != key_id or item.get("role") != "normal": | |
| continue | |
| normalized_item = self._reset_daily_if_needed(dict(item)) | |
| grants = _active_package_grants(normalized_item) | |
| grants.append(grant) | |
| next_item = dict(normalized_item) | |
| next_item["package_grants"] = grants | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| raise ValueError("普通用户不存在") | |
| def add_package_grant(self, key_id: str, package_plan: dict[str, object], *, source: str = "", code_id: str = "") -> dict[str, object]: | |
| normalized_id = self._clean(key_id) | |
| if not normalized_id: | |
| raise ValueError("用户不存在") | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| return self._add_package_grant_in_items( | |
| items, | |
| normalized_id, | |
| package_plan, | |
| source=source, | |
| code_id=code_id, | |
| ) | |
| return self._update_items_locked(mutate) | |
| def reserve_credits(self, key_id: str, amount: int | float, *, model_id: str = "") -> dict[str, object]: | |
| """按固定优先级从各积分池预扣,并在返回值里带上 `_credit_reservation` 明细。 | |
| 扣费顺序(原则:先到期的先用): | |
| 1. 每日免费额度 `daily_credits_remaining` —— 当天有效,次日清零,最先扣 | |
| 2. 套餐券 `package_grants` —— 受日 / 周 / 月限额约束 | |
| 3. 限时券 `credit_grants` —— 带 `expires_at`,先到期的先扣 | |
| 4. 基础积分池 `credits` —— 永久有效,最后扣 | |
| 签到赠送的积分(`perform_checkin` 发放、无 `expires_at`)落在第 4 项基础积分池, | |
| 即在整个链条上排最后;但在**该池内部**它最先被消耗(赠送优先消费, | |
| 见 `consume_reserved_credits` 里的签到账本)。各池的消耗归属见 | |
| `test/test_checkin_credit_ledger.py`。 | |
| """ | |
| amount_decimal = _credit_decimal(amount) | |
| if amount_decimal <= 0: | |
| result = self.credit_summary({"id": key_id, "role": "normal"}) | |
| result["_credit_reservation"] = {"daily": 0, "packages": [], "paid": 0, "permanent_paid": 0, "grants": []} | |
| return result | |
| normalized_id = self._clean(key_id) | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object]]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| normalized_item = self._reset_daily_if_needed(dict(item)) | |
| credits = _credit_decimal(normalized_item.get("credits")) | |
| grants = _normalize_credit_grants(normalized_item.get("credit_grants")) | |
| package_grants = _active_package_grants(normalized_item) | |
| expiring_available = sum((_credit_decimal(grant.get("remaining")) for grant in grants if _is_future_datetime(grant.get("expires_at"))), Decimal("0")) | |
| daily_available = _credit_decimal(self._daily_remaining(normalized_item, model_id)) | |
| daily_to_reserve = min(amount_decimal, daily_available) | |
| remaining_to_reserve = amount_decimal - daily_to_reserve | |
| package_reservations: list[dict[str, object]] = [] | |
| for grant in package_grants: | |
| if remaining_to_reserve <= 0: | |
| break | |
| take = min(_credit_decimal(_package_available(grant)), remaining_to_reserve) | |
| if take <= 0: | |
| continue | |
| if _to_nonnegative_int(grant.get("total_limit")) > 0: | |
| grant["reserved_total"] = _credit_number(_credit_decimal(grant.get("reserved_total")) + take) | |
| if _to_nonnegative_int(grant.get("daily_limit")) > 0: | |
| grant["reserved_daily"] = _credit_number(_credit_decimal(grant.get("reserved_daily")) + take) | |
| if _to_nonnegative_int(grant.get("weekly_limit")) > 0: | |
| grant["reserved_weekly"] = _credit_number(_credit_decimal(grant.get("reserved_weekly")) + take) | |
| if _to_nonnegative_int(grant.get("monthly_limit")) > 0: | |
| grant["reserved_monthly"] = _credit_number(_credit_decimal(grant.get("reserved_monthly")) + take) | |
| package_reservations.append({"id": grant.get("id"), "amount": _credit_number(take)}) | |
| remaining_to_reserve -= take | |
| paid_to_reserve = remaining_to_reserve | |
| if credits + expiring_available < paid_to_reserve: | |
| raise ValueError("剩余积分不足,请先充值") | |
| next_item = dict(normalized_item) | |
| remaining_paid_to_reserve = paid_to_reserve | |
| grant_reservations: list[dict[str, object]] = [] | |
| for grant in grants: | |
| if remaining_paid_to_reserve <= 0 or not _is_future_datetime(grant.get("expires_at")): | |
| continue | |
| grant_remaining = _credit_decimal(grant.get("remaining")) | |
| take = min(grant_remaining, remaining_paid_to_reserve) | |
| if take <= 0: | |
| continue | |
| grant["remaining"] = _credit_number(grant_remaining - take) | |
| grant["reserved"] = _credit_number(_credit_decimal(grant.get("reserved")) + take) | |
| grant_reservations.append({"id": grant.get("id"), "amount": _credit_number(take)}) | |
| remaining_paid_to_reserve -= take | |
| permanent_to_reserve = remaining_paid_to_reserve | |
| next_item["credits"] = _credit_number(credits - permanent_to_reserve) | |
| next_item["reserved_credits"] = _credit_number(_credit_decimal(normalized_item.get("reserved_credits")) + permanent_to_reserve) | |
| next_item["daily_reserved_credits"] = _credit_number(_credit_decimal(normalized_item.get("daily_reserved_credits")) + daily_to_reserve) | |
| next_item["credit_grants"] = grants | |
| next_item["package_grants"] = package_grants | |
| result = self._public_item(next_item) | |
| result["_credit_reservation"] = { | |
| "daily": _credit_number(daily_to_reserve), | |
| "packages": package_reservations, | |
| "paid": _credit_number(paid_to_reserve), | |
| "permanent_paid": _credit_number(permanent_to_reserve), | |
| "grants": grant_reservations, | |
| } | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, result | |
| raise ValueError("普通用户不存在") | |
| return self._update_items_locked(mutate) | |
| def consume_reserved_credits( | |
| self, | |
| key_id: str, | |
| amount: int | float, | |
| *, | |
| daily_amount: int | float | None = None, | |
| package_reservations: list[dict[str, object]] | None = None, | |
| paid_amount: int | float | None = None, | |
| grant_reservations: list[dict[str, object]] | None = None, | |
| ) -> dict[str, object] | None: | |
| amount_decimal = _credit_decimal(amount) | |
| if amount_decimal <= 0: | |
| return self.get_key(key_id, role="normal") | |
| normalized_id = self._clean(key_id) | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object] | None]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| normalized_item = self._reset_daily_if_needed(dict(item)) | |
| daily_reserved = _credit_decimal(normalized_item.get("daily_reserved_credits")) | |
| paid_reserved = _credit_decimal(normalized_item.get("reserved_credits")) | |
| grants = _normalize_credit_grants(normalized_item.get("credit_grants")) | |
| package_grants = _settleable_package_grants(normalized_item) | |
| package_consumed = Decimal("0") | |
| for reservation in package_reservations or []: | |
| if not isinstance(reservation, dict): | |
| continue | |
| package_id = self._clean(reservation.get("id")) | |
| requested = _credit_decimal(reservation.get("amount")) | |
| if not package_id or requested <= 0: | |
| continue | |
| for grant in package_grants: | |
| if self._clean(grant.get("id")) != package_id: | |
| continue | |
| consumed_from_package = Decimal("0") | |
| if _to_nonnegative_int(grant.get("total_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_total"))) | |
| grant["reserved_total"] = _credit_number(_credit_decimal(grant.get("reserved_total")) - take) | |
| grant["used_total"] = _credit_number(_credit_decimal(grant.get("used_total")) + take) | |
| consumed_from_package = max(consumed_from_package, take) | |
| if _to_nonnegative_int(grant.get("daily_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_daily"))) | |
| grant["reserved_daily"] = _credit_number(_credit_decimal(grant.get("reserved_daily")) - take) | |
| grant["used_daily"] = _credit_number(_credit_decimal(grant.get("used_daily")) + take) | |
| consumed_from_package = max(consumed_from_package, take) | |
| if _to_nonnegative_int(grant.get("weekly_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_weekly"))) | |
| grant["reserved_weekly"] = _credit_number(_credit_decimal(grant.get("reserved_weekly")) - take) | |
| grant["used_weekly"] = _credit_number(_credit_decimal(grant.get("used_weekly")) + take) | |
| consumed_from_package = max(consumed_from_package, take) | |
| if _to_nonnegative_int(grant.get("monthly_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_monthly"))) | |
| grant["reserved_monthly"] = _credit_number(_credit_decimal(grant.get("reserved_monthly")) - take) | |
| grant["used_monthly"] = _credit_number(_credit_decimal(grant.get("used_monthly")) + take) | |
| consumed_from_package = max(consumed_from_package, take) | |
| package_consumed += consumed_from_package | |
| break | |
| grant_consumed = Decimal("0") | |
| for reservation in grant_reservations or []: | |
| if not isinstance(reservation, dict): | |
| continue | |
| grant_id = self._clean(reservation.get("id")) | |
| requested = _credit_decimal(reservation.get("amount")) | |
| if not grant_id or requested <= 0: | |
| continue | |
| for grant in grants: | |
| if self._clean(grant.get("id")) != grant_id: | |
| continue | |
| consumed_from_grant = min(requested, _credit_decimal(grant.get("reserved"))) | |
| grant["reserved"] = _credit_number(_credit_decimal(grant.get("reserved")) - consumed_from_grant) | |
| grant_consumed += consumed_from_grant | |
| break | |
| if daily_amount is None and paid_amount is None: | |
| daily_consumed = min(amount_decimal, daily_reserved) | |
| paid_consumed = min(amount_decimal - daily_consumed, paid_reserved) | |
| else: | |
| daily_consumed = min(_credit_decimal(daily_amount), daily_reserved) | |
| paid_consumed = min(_credit_decimal(paid_amount), paid_reserved) | |
| consumed = daily_consumed + paid_consumed | |
| next_item = dict(normalized_item) | |
| # 签到账本按「签到积分最先被消耗」扣减;只有从**基础积分池**扣掉的部分算数 | |
| # (每日免费额度 / 限时券 / 套餐券各自独立,不动签到账本)。 | |
| checkin_remaining = _credit_decimal(normalized_item.get("checkin_credit_remaining")) | |
| if checkin_remaining > 0 and paid_consumed > 0: | |
| checkin_remaining = max(Decimal("0"), checkin_remaining - paid_consumed) | |
| next_item["checkin_credit_remaining"] = _credit_number(checkin_remaining) | |
| next_item["daily_reserved_credits"] = _credit_number(daily_reserved - daily_consumed) | |
| next_item["daily_credits_used"] = _credit_number(_credit_decimal(normalized_item.get("daily_credits_used")) + daily_consumed) | |
| next_item["reserved_credits"] = _credit_number(paid_reserved - paid_consumed) | |
| next_item["credit_grants"] = grants | |
| next_item["package_grants"] = package_grants | |
| next_item["used_credits"] = _credit_number(_credit_decimal(normalized_item.get("used_credits")) + consumed + grant_consumed + package_consumed) | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| return items, None | |
| return self._update_items_locked(mutate) | |
| def refund_reserved_credits( | |
| self, | |
| key_id: str, | |
| amount: int | float, | |
| *, | |
| daily_amount: int | float | None = None, | |
| paid_amount: int | float | None = None, | |
| package_reservations: list[dict[str, object]] | None = None, | |
| grant_reservations: list[dict[str, object]] | None = None, | |
| ) -> dict[str, object] | None: | |
| amount_decimal = _credit_decimal(amount) | |
| if amount_decimal <= 0: | |
| return self.get_key(key_id, role="normal") | |
| normalized_id = self._clean(key_id) | |
| def mutate(items: list[dict[str, object]]) -> tuple[list[dict[str, object]], dict[str, object] | None]: | |
| for index, item in enumerate(items): | |
| if item.get("id") != normalized_id or item.get("role") != "normal": | |
| continue | |
| normalized_item = self._reset_daily_if_needed(dict(item)) | |
| daily_reserved = _credit_decimal(normalized_item.get("daily_reserved_credits")) | |
| paid_reserved = _credit_decimal(normalized_item.get("reserved_credits")) | |
| grants = _normalize_credit_grants(normalized_item.get("credit_grants")) | |
| package_grants = _settleable_package_grants(normalized_item) | |
| for reservation in package_reservations or []: | |
| if not isinstance(reservation, dict): | |
| continue | |
| package_id = self._clean(reservation.get("id")) | |
| requested = _credit_decimal(reservation.get("amount")) | |
| if not package_id or requested <= 0: | |
| continue | |
| for grant in package_grants: | |
| if self._clean(grant.get("id")) != package_id: | |
| continue | |
| if _to_nonnegative_int(grant.get("total_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_total"))) | |
| grant["reserved_total"] = _credit_number(_credit_decimal(grant.get("reserved_total")) - take) | |
| if _to_nonnegative_int(grant.get("daily_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_daily"))) | |
| grant["reserved_daily"] = _credit_number(_credit_decimal(grant.get("reserved_daily")) - take) | |
| if _to_nonnegative_int(grant.get("weekly_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_weekly"))) | |
| grant["reserved_weekly"] = _credit_number(_credit_decimal(grant.get("reserved_weekly")) - take) | |
| if _to_nonnegative_int(grant.get("monthly_limit")) > 0: | |
| take = min(requested, _credit_decimal(grant.get("reserved_monthly"))) | |
| grant["reserved_monthly"] = _credit_number(_credit_decimal(grant.get("reserved_monthly")) - take) | |
| break | |
| for reservation in grant_reservations or []: | |
| if not isinstance(reservation, dict): | |
| continue | |
| grant_id = self._clean(reservation.get("id")) | |
| requested = _credit_decimal(reservation.get("amount")) | |
| if not grant_id or requested <= 0: | |
| continue | |
| for grant in grants: | |
| if self._clean(grant.get("id")) != grant_id: | |
| continue | |
| refunded_from_grant = min(requested, _credit_decimal(grant.get("reserved"))) | |
| grant["reserved"] = _credit_number(_credit_decimal(grant.get("reserved")) - refunded_from_grant) | |
| grant["remaining"] = _credit_number(_credit_decimal(grant.get("remaining")) + refunded_from_grant) | |
| break | |
| if daily_amount is None and paid_amount is None: | |
| daily_refunded = min(amount_decimal, daily_reserved) | |
| paid_refunded = min(amount_decimal - daily_refunded, paid_reserved) | |
| else: | |
| daily_refunded = min(_credit_decimal(daily_amount), daily_reserved) | |
| paid_refunded = min(_credit_decimal(paid_amount), paid_reserved) | |
| next_item = dict(normalized_item) | |
| next_item["daily_reserved_credits"] = _credit_number(daily_reserved - daily_refunded) | |
| next_item["reserved_credits"] = _credit_number(paid_reserved - paid_refunded) | |
| next_item["credits"] = _credit_number(_credit_decimal(normalized_item.get("credits")) + paid_refunded) | |
| next_item["credit_grants"] = grants | |
| next_item["package_grants"] = package_grants | |
| next_items = list(items) | |
| next_items[index] = next_item | |
| return next_items, self._public_item(next_item) | |
| return items, None | |
| return self._update_items_locked(mutate) | |
| def authenticate(self, raw_key: str, *, touch: bool = True) -> dict[str, object] | None: | |
| candidate = self._clean(raw_key) | |
| if not candidate: | |
| return None | |
| candidate_hash = _hash_key(candidate) | |
| def find_match(items: list[dict[str, object]]) -> tuple[int, dict[str, object]] | None: | |
| for index, item in enumerate(items): | |
| if not bool(item.get("enabled", True)): | |
| continue | |
| stored_hash = self._clean(item.get("key_hash")) | |
| if stored_hash and hmac.compare_digest(stored_hash, candidate_hash): | |
| return index, item | |
| return None | |
| # ⚠️⚠️ 这一段跑在**事件循环线程**上(async 路由里同步调 require_identity), | |
| # 所以里面**绝不允许出现任何存储 IO**。锁内只做内存查找 + 内存更新。 | |
| missed = False | |
| public_item: dict[str, object] | None = None | |
| with self._lock: | |
| matched = find_match(self._items) | |
| if matched is None: | |
| # ⚠️ 缓存未命中**也不回读存储**:读库同样要排队等图片任务线程的写锁, | |
| # 而这里跑在事件循环上。单副本部署下内存快照就是权威 | |
| # (创建/修改/删除密钥都会同步更新 self._items), | |
| # 回读只对「外部进程直接改库」有意义 —— 交给后台线程兜底。 | |
| self._reload_requested = True | |
| missed = True | |
| else: | |
| index, matched_item = matched | |
| if touch: | |
| # last_used_at 节流:同一密钥 60 秒内只回写一次。 | |
| throttle_key = self._clean(matched_item.get("id")) or candidate_hash | |
| now_ts = time.time() | |
| if now_ts - self._last_touch_ts.get(throttle_key, 0.0) >= self._TOUCH_INTERVAL_SECONDS: | |
| self._last_touch_ts[throttle_key] = now_ts | |
| if len(self._last_touch_ts) > 4096: | |
| cutoff = now_ts - self._TOUCH_INTERVAL_SECONDS | |
| self._last_touch_ts = { | |
| key: ts for key, ts in self._last_touch_ts.items() if ts >= cutoff | |
| } | |
| # 只更新内存快照 + 记脏,真正的落库交给后台批量 flush。 | |
| touched_at = datetime.now(timezone.utc).isoformat() | |
| next_item = dict(matched_item) | |
| next_item["last_used_at"] = touched_at | |
| self._items[index] = next_item | |
| key_id = self._clean(next_item.get("id")) | |
| if key_id: | |
| self._pending_touch[key_id] = touched_at | |
| matched_item = next_item | |
| public_item = self._public_item(matched_item) | |
| if missed: | |
| # 兜底回读放到后台(每 30 秒一轮),本请求立刻返回 401,不阻塞事件循环。 | |
| self._ensure_touch_flusher() | |
| return None | |
| return public_item | |
| auth_service = AuthService(config.get_user_storage_backend()) | |