diff --git a/app/credential_actions.py b/app/credential_actions.py index fa4c23a..4141b06 100644 --- a/app/credential_actions.py +++ b/app/credential_actions.py @@ -9,11 +9,18 @@ def run(gateway, action, identity=None, *, consent_revision=None): - if action not in {"refresh", "checkin", "sync", "travel", "travel-status", "trial"} or (action in {"refresh", "travel", "travel-status", "trial"} and identity is None): + if action not in {"refresh", "checkin", "sync", "travel", "travel-status", "trial", "reset-cooldown"} or (action in {"refresh", "travel", "travel-status", "trial", "reset-cooldown"} and identity is None): raise HTTPException(404, "凭证操作不存在") if consent_revision is not None and (action != "travel" or identity is None or consent_revision != buddy.AGREEMENT_REVISION): raise HTTPException(400, "首领确认无效或协议版本已变化") config = gateway.CONFIG + # Clearing a cooldown only edits local state: it needs neither the ledger nor a network round + # trip, and it must stay available while the periodic maintenance sweep holds its lock. It is + # also allowed for a disabled account, since a stale cooldown is worth clearing either way. + if action == "reset-cooldown": + result = _reset_cooldowns(gateway, identity) + _audit(config, result) + return {"ok": result["ok"], "results": [result]} pool, ledger = config.get("cred_pool"), config.get("ledger") if pool is None or (action not in {"refresh", "trial"} and (ledger is None or gateway.credits_mod is None)): raise HTTPException(503, "凭证维护尚未就绪") @@ -49,24 +56,7 @@ def run(gateway, action, identity=None, *, consent_revision=None): result["checkin_ok"] = True result.update(ok=False, message="操作失败，保留已有数据；请检查账号状态后重试") results.append(result) - audit = config.get("audit_store") - if audit: - try: - details = {"credential": result["id"], "ok": result["ok"]} - if action == "trial": - details.update(outcome="success" if result["ok"] else "error", stage=result.get("state"), - status_code=result.get("status"), - code=str(result["code"]) if result.get("code") is not None else None, - duration_ms=(time.monotonic() - started) * 1000) - trip = result.get("travel") if action == "checkin" else result if action in {"travel", "travel-status"} else None - if isinstance(trip, dict): - details.update(outcome="success" if trip.get("ok") else "warning" if trip.get("buddy_blocked") else "error", - stage=trip.get("phase"), status_code=trip.get("http_status"), - code=trip.get("reason") or (str(trip["code"]) if trip.get("code") is not None else None), - consent_source=trip.get("consent_source")) - audit.event("admin", "credential." + action, details) - except Exception: - pass + _audit(config, result, started) response = {"ok": bool(results) and all(r["ok"] for r in results), "results": results} if identity is None and ledger is not None: response["credits"] = ledger.snapshot() # Keep the legacy check-in response field. @@ -75,6 +65,57 @@ def run(gateway, action, identity=None, *, consent_revision=None): gateway._HOUSEKEEP_LOCK.release() +def _audit(config, result, started=None): + """Record one credential action; auditing must never break the action itself.""" + audit = config.get("audit_store") + if not audit: + return + try: + action = result.get("action") + details = {"credential": result["id"], "ok": result["ok"]} + if action == "trial": + details.update(outcome="success" if result["ok"] else "error", stage=result.get("state"), + status_code=result.get("status"), + code=str(result["code"]) if result.get("code") is not None else None, + duration_ms=(time.monotonic() - started) * 1000 if started is not None else None) + if action == "reset-cooldown": + # The audit sanitizer keeps a fixed key allowlist, so report the split outcome through + # fields it retains rather than widening a shared schema. + details.update(outcome="success" if result["ok"] else "error", + stage="durable" if result.get("durable") else "memory_only", + code="reset_cooldown" if result.get("durable") else "reset_cooldown_write_failed") + trip = result.get("travel") if action == "checkin" else result if action in {"travel", "travel-status"} else None + if isinstance(trip, dict): + details.update(outcome="success" if trip.get("ok") else "warning" if trip.get("buddy_blocked") else "error", + stage=trip.get("phase"), status_code=trip.get("http_status"), + code=trip.get("reason") or (str(trip["code"]) if trip.get("code") is not None else None), + consent_source=trip.get("consent_source")) + audit.event("admin", "credential." + str(action), details) + except Exception: + pass + + +def _reset_cooldowns(gateway, identity): + """Lift every cooldown an account holds, reporting memory and durability separately.""" + pool = gateway.CONFIG.get("cred_pool") + if pool is None: + raise HTTPException(503, "凭证维护尚未就绪") + pool._rescan() + entry = next((e for e in pool.entries() if e.get("account_key") == identity), None) + if entry is None: + raise HTTPException(404, "凭证不存在或身份已变化") + outcome = pool.reset_cooldowns_for(identity) + result = {"id": identity, "name": Path(entry["id"]).name, "action": "reset-cooldown", + "ok": outcome["durable"], "changed_in_memory": outcome["changed_in_memory"], + "durable": outcome["durable"]} + if outcome["durable"]: + result["message"] = "已清除该账号的冷却" if outcome["changed_in_memory"] else "该账号当前没有冷却" + else: + # Never let a failed write read as a completed reset: a restart would restore the row. + result["message"] = "内存冷却已清除，但写入失败，重启后可能恢复；请稍后重试" + return result + + def _one(gateway, pool, ledger, entry, action, *, automatic=False, consent_revision=None): if action == "trial": return trial_management.perform(gateway, pool, entry) diff --git a/app/credential_cooldowns.py b/app/credential_cooldowns.py new file mode 100644 index 0000000..bac7f40 --- /dev/null +++ b/app/credential_cooldowns.py @@ -0,0 +1,435 @@ +#!/usr/bin/env python3 +"""Persist credential circuit-breaker and per-model 429 cooldowns across restarts. + +Cooldowns are advisory: they only delay the next attempt against a backend that just +refused, so losing one costs a single wasted request. This store is therefore tolerant — +unusable storage degrades to the previous in-memory behaviour and never disables an +account. The properties that must hold on every path are: + +* a restored deadline can never sit further out than its ceiling, so a stale or crafted + file cannot park an account permanently; +* what the writer produces is always something the reader accepts, because a file the + reader rejects discards *every* cooldown it held, not just the offending row. +""" + +from __future__ import annotations + +import json +import os +import re +import stat +import tempfile +import threading +import time + +VERSION = 1 +MAX_ACCOUNTS = 256 # Bounded table; matches the credential pool's practical size. +MAX_MODELS = 64 # Bounded per-account model rows. +MAX_BYTES = 256 * 1024 # Bounded read *and* write; the writer never exceeds this. +# A credential breaker is short-lived (CRED_COOLDOWN is 300s), so it gets a much tighter +# ceiling than a per-model quota cooldown, which may legitimately run to MODEL_COOLDOWN_MAX. +AUTH_CEILING_S = 3600 +MODEL_CEILING_S = 24 * 3600 + +_IDENTITY = re.compile(r"[0-9a-f]{64}") +# Model and profile identifiers may contain slashes (vendor-qualified names), so allow the +# punctuation seen in routing rather than assuming a bare token. +_TOKEN = re.compile(r"[A-Za-z0-9_.:/\-]{1,128}") +_REASON_CHARS = 256 +_FIELDS = {"profile", "fail_until", "reason", "failed_at", "models"} + + +def _text(value, limit=_REASON_CHARS): + """Accept a bounded printable string; never raise on hostile input.""" + if not isinstance(value, str): + return "" + return "".join(ch for ch in value if ch.isprintable())[:limit] + + +def _number(value): + """Accept only a finite number; bool is not a number, and a huge int must not raise.""" + if type(value) not in (int, float): + return None + try: + number = float(value) + except (OverflowError, ValueError): + return None + if number != number or number in (float("inf"), float("-inf")): + return None + return number + + +def _deadline(value, now, ceiling): + """Accept only a finite deadline still in the future and within the ceiling.""" + number = _number(value) + if number is None or number <= now or number > now + ceiling: + return None + return number + + +def _valid_identity(value): + return isinstance(value, str) and _IDENTITY.fullmatch(value) is not None + + +def _valid_token(value): + return isinstance(value, str) and _TOKEN.fullmatch(value) is not None + + +def _strict_json(raw: bytes): + """Parse JSON while rejecting duplicate keys and Infinity/NaN.""" + def no_duplicates(pairs): + seen = {} + for key, value in pairs: + if key in seen: + raise ValueError("duplicate key") + seen[key] = value + return seen + + def no_constants(name): + raise ValueError(name) + + return json.loads(raw.decode("utf-8"), object_pairs_hook=no_duplicates, parse_constant=no_constants) + + +class CredentialCooldowns: + """Thread-safe, optionally persisted cooldown table keyed by validated account identity.""" + + def __init__(self, path=None, *, auth_ceiling_s: float = AUTH_CEILING_S, + model_ceiling_s: float = MODEL_CEILING_S): + self.path = str(path) if path else None + self.auth_ceiling_s = _bounded_ceiling(auth_ceiling_s, AUTH_CEILING_S) + self.model_ceiling_s = _bounded_ceiling(model_ceiling_s, MODEL_CEILING_S) + self._lock = threading.RLock() + self._data: dict[str, dict] = {} + # Last write failure, surfaced for diagnostics instead of failing silently. + self.last_error: str | None = None + # Set when the in-memory table changed but the write did not land. A later clear must + # retry the write instead of reporting "nothing to do" while stale rows sit on disk. + self._dirty = False + if self.path: + self._load() + + # -- persistence ------------------------------------------------------- + + def _load(self): + """Adopt only a fully valid snapshot; anything else leaves the table empty.""" + try: + # Reject a symlink or FIFO *before* opening, so a device cannot block the read. + # This is best effort on Windows and is backed up by the fstat check below. + if not stat.S_ISREG(os.lstat(self.path).st_mode): + return + fd = os.open(self.path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + | getattr(os, "O_NONBLOCK", 0)) + except OSError: + return + try: + with os.fdopen(fd, "rb") as stream: + if not stat.S_ISREG(os.fstat(stream.fileno()).st_mode): + return + raw = stream.read(MAX_BYTES + 1) + except OSError: + return + if len(raw) > MAX_BYTES: + return + try: + document = _strict_json(raw) + except (ValueError, UnicodeDecodeError, RecursionError): + return + if (not isinstance(document, dict) or set(document) != {"version", "accounts"} + or type(document["version"]) is not int or document["version"] != VERSION): + return + accounts = document["accounts"] + if not isinstance(accounts, dict) or len(accounts) > MAX_ACCOUNTS: + return + now = time.time() + restored: dict[str, dict] = {} + for identity, row in accounts.items(): + if not _valid_identity(identity) or not isinstance(row, dict) or set(row) != _FIELDS: + return + profile, models = row["profile"], row["models"] + if not _valid_token(profile) or not isinstance(models, dict) or len(models) > MAX_MODELS: + return + fail_until = _deadline(row["fail_until"], now, self.auth_ceiling_s) + kept = {} + for model, until in models.items(): + if not _valid_token(model): + return + deadline = _deadline(until, now, self.model_ceiling_s) + if deadline is not None: + kept[model] = deadline + if fail_until is None and not kept: + continue + restored[identity] = {"profile": profile, "fail_until": fail_until or 0.0, + "reason": _text(row["reason"]), + "failed_at": _number(row["failed_at"]) or 0.0, "models": kept} + with self._lock: + self._data = restored + + def _serialize_locked(self): + """Build the exact payload to write, shedding rows until it fits MAX_BYTES. + + Returns the payload plus whether rows had to be shed, because a caller that asked to + record a cooldown must not be told its write was durable when that row was the one + given up to satisfy the byte bound. + """ + accounts = dict(self._data) + shed = False + while True: + try: + content = json.dumps({"version": VERSION, "accounts": accounts}, + ensure_ascii=False, separators=(",", ":"), + allow_nan=False).encode("utf-8") + except (TypeError, ValueError): + return None, shed + if len(content) <= MAX_BYTES or not accounts: + return content, shed + # Shed the accounts nearest expiry first, in one proportional step. + ordered = sorted(accounts, key=lambda key: max( + [float(accounts[key].get("fail_until") or 0.0)] + list(accounts[key]["models"].values()))) + keep = max(0, int(len(ordered) * MAX_BYTES / len(content) * 0.9)) + for key in ordered[:max(1, len(ordered) - keep)]: + accounts.pop(key, None) + shed = True + + def _save_locked(self) -> bool: + """Atomically rewrite the table; returns False when the update was not fully durable. + + The payload is written even when rows had to be shed, so the file stays bounded and + readable; the False return tells the caller its update may not have survived. + """ + if not self.path: + return True + content, shed = self._serialize_locked() + if content is None: + self.last_error = "serialize" + self._dirty = True + return False + temporary = None + try: + directory = os.path.dirname(self.path) or "." + os.makedirs(directory, mode=0o700, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=".credential-cooldowns-", suffix=".tmp", dir=directory) + with os.fdopen(fd, "wb") as stream: + stream.write(content) + try: + os.chmod(temporary, 0o600) + except OSError: + pass # chmod does not establish owner-only ACLs on Windows. + os.replace(temporary, self.path) + except OSError as error: + self.last_error = type(error).__name__ + self._dirty = True + return False + finally: + if temporary: + try: + os.unlink(temporary) + except OSError: + pass + self.last_error = "capacity" if shed else None + self._dirty = shed # The write landed, but rows were dropped to fit. + return not shed + + # -- table helpers ----------------------------------------------------- + + @staticmethod + def _blank(profile): + return {"profile": profile, "fail_until": 0.0, "reason": "", "failed_at": 0.0, "models": {}} + + def _lookup_locked(self, identity, profile): + """Read-only lookup; a row owned by another product is not visible.""" + if not _valid_identity(identity): + return None + row = self._data.get(identity) + if row is None or (profile is not None and row.get("profile") != profile): + return None + return row + + def _ensure_locked(self, identity, profile): + """Write-side lookup that adopts a new product identity for a reused path. + + Rejects values the reader would refuse, because writing one unreadable row would + make the reader discard the whole file on the next start. + """ + if not _valid_identity(identity) or not _valid_token(profile): + return None + row = self._data.get(identity) + if row is None: + if len(self._data) >= MAX_ACCOUNTS: + self._drop_least_useful_locked() + row = self._data[identity] = self._blank(profile) + elif row.get("profile") != profile: + # A path reused by another product must not inherit the old account's cooldowns. + row = self._data[identity] = self._blank(profile) + return row + + def _drop_least_useful_locked(self): + """Evict the account whose furthest deadline is nearest, to stay within capacity.""" + if not self._data: + return + oldest = min(self._data, key=lambda key: max( + [float(self._data[key].get("fail_until") or 0.0)] + list(self._data[key]["models"].values()))) + self._data.pop(oldest, None) + + def _retry_locked(self) -> bool: + """Re-attempt a write that previously failed, so stale disk rows are not left behind.""" + if not self._dirty: + return True + return self._save_locked() + + def _drop_locked(self, identity, row): + if not row.get("models") and not row.get("fail_until"): + self._data.pop(identity, None) + + # -- writes ------------------------------------------------------------ + + def note_credential(self, identity, profile, until, reason="", now=None) -> bool: + """Record a credential-wide circuit-breaker deadline; returns whether it is durable.""" + deadline = _number(until) + stamp = _number(time.time() if now is None else now) + if deadline is None or stamp is None: + return False # Reject before mutating anything. + with self._lock: + self._expire_locked(stamp) # Capacity must reflect live rows only. + row = self._ensure_locked(identity, profile) + if row is None: + return False + # Never shorten a live deadline, and never extend past the ceiling. + row["fail_until"] = max(float(row.get("fail_until") or 0.0), + min(deadline, stamp + self.auth_ceiling_s)) + row["reason"] = _text(reason) + row["failed_at"] = stamp + return self._save_locked() + + def note_model(self, identity, profile, model, until, now=None) -> bool: + """Record a per-model 429 deadline; returns whether it is durable.""" + deadline = _number(until) + stamp = _number(time.time() if now is None else now) + if not _valid_token(model) or deadline is None or stamp is None: + return False # Reject before mutating anything. + with self._lock: + self._expire_locked(stamp) # Capacity must reflect live rows only. + row = self._ensure_locked(identity, profile) + if row is None: + return False + models = row["models"] + if model not in models and len(models) >= MAX_MODELS: + # Evict the model cooldown nearest expiry rather than refusing the new one. + models.pop(min(models, key=lambda name: models[name]), None) + models[model] = max(float(models.get(model) or 0.0), + min(deadline, stamp + self.model_ceiling_s)) + return self._save_locked() + + def _expire_locked(self, now): + """Drop expired deadlines so capacity reflects live rows only.""" + for identity in list(self._data): + row = self._data[identity] + if row.get("fail_until") and float(row["fail_until"]) <= now: + row["fail_until"] = 0.0 + row["models"] = {m: u for m, u in row["models"].items() if u > now} + self._drop_locked(identity, row) + + def clear_credential(self, identity, profile=None) -> dict: + """Drop a circuit breaker, reporting whether anything changed and whether it is durable. + + "Nothing to change" and "the write failed" are different outcomes: only the second + leaves stale data on disk that a later restart would restore. + """ + with self._lock: + row = self._lookup_locked(identity, profile) + if row is None: + # Nothing to change here, but a previous failed write may still be on disk. + return {"changed": False, "durable": self._retry_locked()} + row["fail_until"] = 0.0 + row["reason"] = "" + self._drop_locked(identity, row) + return {"changed": True, "durable": self._save_locked()} + + def clear_model(self, identity, profile, model) -> dict: + """Drop one model's cooldown, reporting whether anything changed and whether it is durable.""" + with self._lock: + row = self._lookup_locked(identity, profile) + if row is None or model not in row["models"]: + return {"changed": False, "durable": self._retry_locked()} + row["models"].pop(model, None) + self._drop_locked(identity, row) + return {"changed": True, "durable": self._save_locked()} + + def forget(self, identity) -> dict: + """Remove an account entirely; used when its credential is deleted or replaced.""" + with self._lock: + if self._data.pop(identity, None) is None: + return {"changed": False, "durable": self._retry_locked()} + return {"changed": True, "durable": self._save_locked()} + + def prune(self, now=None) -> dict: + """Drop expired deadlines so the table cannot accumulate dead rows.""" + now = time.time() if now is None else now + with self._lock: + changed = False + for identity in list(self._data): + row = self._data[identity] + if row.get("fail_until") and float(row["fail_until"]) <= now: + row["fail_until"] = 0.0 + changed = True + kept = {m: u for m, u in row["models"].items() if u > now} + if kept != row["models"]: + row["models"] = kept + changed = True + before = len(self._data) + self._drop_locked(identity, row) + changed = changed or len(self._data) != before + return {"changed": changed, "durable": self._save_locked() if changed else True} + + # -- reads ------------------------------------------------------------- + + def credential_until(self, identity, profile=None) -> float: + """Return a live circuit-breaker deadline, or zero.""" + now = time.time() + with self._lock: + row = self._lookup_locked(identity, profile) + if row is None: + return 0.0 + until = float(row.get("fail_until") or 0.0) + return until if until > now else 0.0 + + def model_until(self, identity, profile, model) -> float: + """Return a live per-model cooldown deadline, or zero.""" + now = time.time() + with self._lock: + row = self._lookup_locked(identity, profile) + if row is None: + return 0.0 + until = float(row["models"].get(model) or 0.0) + return until if until > now else 0.0 + + def restore(self, identity, profile) -> dict: + """Return the persisted state for one account as a plain dict.""" + now = time.time() + with self._lock: + row = self._lookup_locked(identity, profile) + if row is None: + return {} + until = float(row.get("fail_until") or 0.0) + return {"fail_until": until if until > now else 0.0, + "reason": row.get("reason") or "", + "failed_at": float(row.get("failed_at") or 0.0), + "models": {m: u for m, u in row["models"].items() if u > now}} + + def detail(self) -> list: + """Return a bounded snapshot for diagnostics.""" + now = time.time() + with self._lock: + return [{"identity": identity, "profile": row.get("profile"), + "fail_until": float(row.get("fail_until") or 0.0), + "reason": row.get("reason") or "", + "models": {m: u for m, u in row["models"].items() if u > now}} + for identity, row in self._data.items()] + + +def _bounded_ceiling(value, default): + """Keep a caller-supplied ceiling finite and within the module default.""" + number = _number(value) + if number is None or number < 60: + return float(default) + return min(number, float(default)) diff --git a/app/usage_snapshots.py b/app/usage_snapshots.py new file mode 100644 index 0000000..b4f840d --- /dev/null +++ b/app/usage_snapshots.py @@ -0,0 +1,329 @@ +#!/usr/bin/env python3 +"""Cache per-account usage snapshots so the dashboard is not blank after a restart. + +Usage is presentation data: it feeds the billing and dashboard totals and is never +consulted when choosing a credential, so a cache miss or a corrupt file costs only a +cosmetic gap until the next refresh. That is exactly why this store must never become a +source of truth — it is rebuilt from upstream, and a cached snapshot is always treated as +stale-but-usable rather than authoritative. + +The file is keyed by credential path, matching how the aggregate is assembled, but each +row also records the account identity it was fetched for. A path reused by a different +account therefore cannot inherit the previous account's usage. +""" + +from __future__ import annotations + +import json +import os +import re +import stat +import tempfile +import threading +import time + +VERSION = 1 +MAX_ACCOUNTS = 256 # Bounded table; matches the credential pool's practical size. +MAX_DAYS = 31 # The upstream usage window is 30 days; allow one spare day. +MAX_MODELS = 128 # Bounded model rows per day. +MAX_BYTES = 512 * 1024 # Bounded read *and* write; the writer never exceeds this. +MAX_CREDITS = 1e9 # A single day/model credit figure is far below this. + +_IDENTITY = re.compile(r"[0-9a-f]{64}") +_DATE = re.compile(r"\d{4}-\d{2}-\d{2}") +_FIELDS = {"identity", "site", "by_day", "total_credits", "requests", "partial", "fetched_at"} + + +def _number(value, *, low=None, high=None): + """Accept only a finite number inside the given range; bool is not a number.""" + if type(value) not in (int, float): + return None + try: + number = float(value) + except (OverflowError, ValueError): + return None + if number != number or number in (float("inf"), float("-inf")): + return None + if low is not None and number < low: + return None + if high is not None and number > high: + return None + return number + + +def _valid_date(value): + """Accept a real calendar date; a shape-only match would allow 2026-02-31.""" + if _DATE.fullmatch(value) is None: + return False + try: + time.strptime(value, "%Y-%m-%d") + except ValueError: + return False + return True + + +def _valid_identity(value): + return isinstance(value, str) and _IDENTITY.fullmatch(value) is not None + + +def _valid_site(value): + return value in ("domestic", "international") + + +def _strict_json(raw: bytes): + """Parse JSON while rejecting duplicate keys and Infinity/NaN.""" + def no_duplicates(pairs): + seen = {} + for key, value in pairs: + if key in seen: + raise ValueError("duplicate key") + seen[key] = value + return seen + + def no_constants(name): + raise ValueError(name) + + return json.loads(raw.decode("utf-8"), object_pairs_hook=no_duplicates, parse_constant=no_constants) + + +class UsageSnapshots: + """Thread-safe, optionally persisted cache of per-account usage snapshots.""" + + def __init__(self, path=None): + self.path = str(path) if path else None + self._lock = threading.RLock() + self._data: dict[str, dict] = {} + self.last_error: str | None = None + # Set when the cache changed but the write did not land, so a later forget retries. + self._dirty = False + if self.path: + self._load() + + # -- persistence ------------------------------------------------------- + + def _load(self): + """Adopt only a fully valid snapshot; anything else leaves the cache empty.""" + try: + if not stat.S_ISREG(os.lstat(self.path).st_mode): + return + fd = os.open(self.path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + | getattr(os, "O_NONBLOCK", 0)) + except OSError: + return + try: + with os.fdopen(fd, "rb") as stream: + if not stat.S_ISREG(os.fstat(stream.fileno()).st_mode): + return + raw = stream.read(MAX_BYTES + 1) + except OSError: + return + if len(raw) > MAX_BYTES: + return + try: + document = _strict_json(raw) + except (ValueError, UnicodeDecodeError, RecursionError): + return + if (not isinstance(document, dict) or set(document) != {"version", "accounts"} + or type(document["version"]) is not int or document["version"] != VERSION): + return + accounts = document["accounts"] + if not isinstance(accounts, dict) or len(accounts) > MAX_ACCOUNTS: + return + restored: dict[str, dict] = {} + for path, row in accounts.items(): + if not isinstance(path, str) or not path or len(path) > 4096: + return + if not isinstance(row, dict) or set(row) != _FIELDS: + return + clean = self._valid_row(row) + if clean is None: + return + restored[path] = clean + with self._lock: + self._data = restored + + @staticmethod + def _valid_row(row): + """Return a normalized row, or None when any field is unusable. + + Malformed values are rejected rather than coerced: turning a bad credit figure into + zero would silently display a wrong number as if it were measured. + """ + if not _valid_identity(row["identity"]) or not _valid_site(row["site"]): + return None + total = _number(row["total_credits"], low=0.0, high=MAX_CREDITS) + requests = _number(row["requests"], low=0.0, high=MAX_CREDITS) + fetched_at = _number(row["fetched_at"], low=0.0) + if total is None or requests is None or fetched_at is None or type(row["partial"]) is not bool: + return None + if requests != int(requests): # A request count is a whole number. + return None + if fetched_at > time.time() + 86400: # A snapshot cannot be fetched in the future. + return None + by_day = row["by_day"] + if not isinstance(by_day, dict) or len(by_day) > MAX_DAYS: + return None + days = {} + for day, models in by_day.items(): + if not isinstance(day, str) or not _valid_date(day) or not isinstance(models, dict): + return None + if len(models) > MAX_MODELS: + return None + clean_models = {} + for model, credit in models.items(): + value = _number(credit, low=0.0, high=MAX_CREDITS) + if not isinstance(model, str) or not model or len(model) > 128 or value is None: + return None + clean_models[model] = value + days[day] = clean_models + return {"identity": row["identity"], "site": row["site"], "by_day": days, + "total_credits": total, "requests": int(requests), + "partial": row["partial"], "fetched_at": fetched_at} + + def _serialize_locked(self): + """Build the payload, shedding the oldest snapshots until it fits MAX_BYTES. + + Returns the payload plus whether rows had to be shed, so a caller is not told its + snapshot was cached when that row was the one given up to satisfy the byte bound. + """ + accounts = dict(self._data) + shed = False + while True: + try: + content = json.dumps({"version": VERSION, "accounts": accounts}, + ensure_ascii=False, separators=(",", ":"), + allow_nan=False).encode("utf-8") + except (TypeError, ValueError): + return None, shed + if len(content) <= MAX_BYTES or not accounts: + return content, shed + ordered = sorted(accounts, key=lambda key: accounts[key].get("fetched_at") or 0.0) + keep = max(0, int(len(ordered) * MAX_BYTES / len(content) * 0.9)) + for key in ordered[:max(1, len(ordered) - keep)]: + accounts.pop(key, None) + shed = True + + def _save_locked(self) -> bool: + """Atomically rewrite the cache; returns False when the update was not fully durable.""" + if not self.path: + return True + content, shed = self._serialize_locked() + if content is None: + self.last_error = "serialize" + self._dirty = True + return False + temporary = None + try: + directory = os.path.dirname(self.path) or "." + os.makedirs(directory, mode=0o700, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=".usage-snapshots-", suffix=".tmp", dir=directory) + with os.fdopen(fd, "wb") as stream: + stream.write(content) + try: + os.chmod(temporary, 0o600) + except OSError: + pass # chmod does not establish owner-only ACLs on Windows. + os.replace(temporary, self.path) + except OSError as error: + self.last_error = type(error).__name__ + self._dirty = True + return False + finally: + if temporary: + try: + os.unlink(temporary) + except OSError: + pass + self.last_error = "capacity" if shed else None + self._dirty = shed + return not shed + + # -- reads / writes ---------------------------------------------------- + + def accounts(self) -> dict: + """Return a copy of the cached snapshots, for merging into the live aggregate.""" + with self._lock: + return {path: dict(row, by_day={day: dict(models) for day, models in row["by_day"].items()}) + for path, row in self._data.items()} + + def store(self, path, identity, site, usage, *, partial=False, now=None) -> bool: + """Cache one account's usage snapshot; returns whether it is durable. + + Malformed input is rejected rather than normalized: turning a missing or bad field + into a plausible zero would present unmeasured data as if it had been measured. + """ + stamp = _number(time.time() if now is None else now, low=0.0) + if not isinstance(path, str) or not path or len(path) > 4096 or stamp is None: + return False + if not isinstance(usage, dict) or type(partial) is not bool: + return False + row = self._valid_row({"identity": identity, "site": site, + "by_day": usage.get("by_day"), + "total_credits": usage.get("total_credits"), + "requests": usage.get("requests"), + "partial": partial, "fetched_at": stamp}) + if row is None: + return False + with self._lock: + if path not in self._data and len(self._data) >= MAX_ACCOUNTS: + self._evict_locked() + # Keep the validated copy so later caller mutation cannot corrupt the cache. + self._data[path] = row + return self._save_locked() + + def forget(self, path) -> dict: + """Drop a deleted credential's snapshot, reporting change and durability separately.""" + with self._lock: + if self._data.pop(path, None) is None: + return {"changed": False, "durable": self._retry_locked()} + return {"changed": True, "durable": self._save_locked()} + + def identity_matches(self, path, identity) -> bool: + """Whether a cached row belongs to the account currently occupying this path.""" + with self._lock: + row = self._data.get(path) + return row is not None and row.get("identity") == identity + + def drop_mismatched(self, pool) -> bool: + """Discard snapshots whose path is now owned by a different account.""" + # Read the pool first: taking the pool lock while holding this store's lock would + # invert the pool -> store order used by store(), which can deadlock. + owners = [(entry["id"], entry.get("account_key")) for entry in pool.entries()] + with self._lock: + dropped = False + for path, identity in owners: + row = self._data.get(path) + if row is not None and identity and row.get("identity") != identity: + self._data.pop(path, None) + dropped = True + if dropped: + self._save_locked() + return dropped + + def _retry_locked(self) -> bool: + """Re-attempt a write that previously failed, so stale disk rows are not left behind.""" + return self._save_locked() if self._dirty else True + + def _evict_locked(self): + if self._data: + self._data.pop(min(self._data, key=lambda key: self._data[key].get("fetched_at") or 0.0), None) + + def prune(self, max_age_s: float = 7 * 24 * 3600, now=None) -> bool: + """Drop snapshots too old to be worth showing; the next refresh replaces them.""" + now = time.time() if now is None else now + with self._lock: + before = len(self._data) + self._data = {path: row for path, row in self._data.items() + if now - float(row.get("fetched_at") or 0.0) <= max_age_s} + if len(self._data) != before: + self._save_locked() + return True + return False + + def detail(self) -> list: + """Return a bounded snapshot for diagnostics.""" + with self._lock: + return [{"path": path, "identity": row["identity"], "site": row["site"], + "total_credits": row["total_credits"], "requests": row["requests"], + "partial": row["partial"], "fetched_at": row["fetched_at"]} + for path, row in self._data.items()] diff --git a/converter.py b/converter.py index 7581585..36198c5 100644 --- a/converter.py +++ b/converter.py @@ -50,7 +50,9 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, from app import auth_oauth from app import trial_rewards from app import buddy, checkin as checkin_service, model_policy, travel +from app.credential_cooldowns import CredentialCooldowns from app.model_blocks import ModelBlocks +from app.usage_snapshots import UsageSnapshots from app.client_hangup import ClientHungUp, await_or_hangup from app.observability import (AuditMiddleware, observe_recovery, observe_route, observe_usage, observe_attempt, observe_failure, @@ -329,6 +331,8 @@ def summary(self) -> dict: } +STORAGE_WARN_INTERVAL = 300 # Rate limit for persistence-failure warnings +USAGE_CACHE_MAX_AGE_S = 7 * 24 * 3600 # A cached usage snapshot older than this is not shown. STICKY_TTL = 30 * 60 # Idle session binding lifetime in seconds STICKY_MAX = 512 # Session binding capacity CRED_COOLDOWN = 300 # Credential cooldown in seconds @@ -461,13 +465,16 @@ class CredentialPool: """Manage credential discovery, reloads, sticky sessions, cooldowns and refresh.""" def __init__(self, paths: list[Path] | None = None, scan: bool = False, - blocks_path: Path | None = None): + blocks_path: Path | None = None, cooldowns_path: Path | None = None): self._lock = threading.RLock() self._entries: list[dict] = [] # {id, cm, fail_until} self._sticky: "OrderedDict[str, tuple[str, float]]" = OrderedDict() self._model_fail: dict[tuple[str, str], float] = {} # Per-credential/model 429 expiry # Keep unsupported-model backoff isolated by backend and model. self._blocks = ModelBlocks(blocks_path, ttl_s=MODEL_SITE_BLOCK_S, max_ttl_s=MODEL_SITE_BLOCK_MAX_S) + # Cooldowns outlive a restart so a backend that just refused is not retried immediately. + self._cooldowns = CredentialCooldowns(cooldowns_path) + self._storage_warned = 0.0 # Rate limit for persistence-failure warnings self._rr = {None: 0, "cn": 0, "intl": 0} self._ledger = None # Prefer credits expiring sooner. self._capacity = AccountCapacity() @@ -505,16 +512,25 @@ def reload(self, paths: list[Path], *, reset: bool = True): changed = reset or generation != entry.get("generation") if changed: old_identity = entry.get("account_key") - if old_identity != identity: + replaced = old_identity != identity + if replaced: self._model_fail = {key: until for key, until in self._model_fail.items() if key[0] != cid} self._sticky = OrderedDict((key, value) for key, value in self._sticky.items() if value[0] != cid) + # The path now belongs to another account; its cooldowns must not carry over. + self.forget_credential_state(entry) if entry.get("uid"): have_uids.pop(old_identity, None) entry.update(uid=summary.get("uid"), profile=summary["profile"], site=summary["site"], account_key=identity, generation=generation, catalog_dirty=True) self._bind_entry(entry) - if reset or old_identity != identity: + if reset or replaced: entry.update(fail_until=0.0, keepalive_after=0.0) + if reset and not replaced: + # An explicit reload re-evaluates auth for the same account, so the + # persisted breaker is lifted; a replacement must keep the incoming + # account's own breaker, which is hydrated just below. + self._forget_credential_cooldown(entry) + self._adopt_cooldowns(entry) if entry.get("uid"): have_uids[identity] = cid self._queue_sync(cid) @@ -536,6 +552,7 @@ def reload(self, paths: list[Path], *, reset: bool = True): "site": summary.get("site"), "profile": profile, "generation": manager._generation, "account_key": identity_key, "catalog_dirty": True} self._bind_entry(entry) + self._adopt_cooldowns(entry) self._entries.append(entry) by_id[cid] = entry self._ignored_duplicates.discard(cid) @@ -618,6 +635,7 @@ def prune(self): for entry in removed: if self._ledger is not None: self._ledger.remove(entry["id"]) + self.forget_credential_state(entry) self._entries = [e for e in self._entries if e not in removed] if len(self._entries) != before: ids = {e["id"] for e in self._entries} @@ -655,6 +673,145 @@ def _bind_entry(self, entry): else: self._ledger.remove(entry["id"]) + def _adopt_cooldowns(self, entry): + """Hydrate persisted cooldowns once per identity, so the in-memory table stays authoritative.""" + identity, profile = entry.get("account_key"), entry.get("profile") + if not self._durable_identity(entry): + return # Adopt later, once this account's identity is validated. + if entry.get("cooldowns_adopted") == identity: + return + entry["cooldowns_adopted"] = identity + state = self._cooldowns.restore(identity, profile) + if not state: + return + # Deadlines are absolute and already bounded, so adopting one never extends a cooldown. + if state.get("fail_until"): + entry["fail_until"] = max(entry.get("fail_until") or 0.0, state["fail_until"]) + if state.get("reason"): + entry["last_error"] = state["reason"] + if state.get("failed_at"): + entry["last_failure_at"] = state["failed_at"] + for model, until in (state.get("models") or {}).items(): + key = (entry["id"], model) + self._model_fail[key] = max(self._model_fail.get(key, 0.0), until) + + def _durable_identity(self, entry) -> bool: + """Whether this account's identity is complete enough to key durable state. + + A hash is derived from the profile and UID, so an account with no UID still yields a + stable-looking hash shared by every other account in the same state. Durable rows + must therefore be keyed only when the identifying components are actually present. + """ + return bool(entry.get("account_key") and entry.get("profile") and entry.get("uid")) + + def _remember_credential(self, entry, reason): + """Mirror a credential circuit breaker to disk; returns whether it is durable.""" + if not self._durable_identity(entry): + return False + durable = self._cooldowns.note_credential(entry["account_key"], entry["profile"], + entry["fail_until"], reason=reason) + self._warn_storage("cooldown", durable) + return durable + + def _remember_model(self, entry, model, until): + """Mirror a model cooldown to disk; returns whether it is durable.""" + if not self._durable_identity(entry): + return False + durable = self._cooldowns.note_model(entry["account_key"], entry["profile"], model, until) + self._warn_storage("cooldown", durable) + return durable + + def _warn_storage(self, label, durable): + """Report a persistence failure operationally, rate limited so a hot path cannot flood.""" + if durable: + self._storage_warned = 0.0 + return + now = time.time() + if now - getattr(self, "_storage_warned", 0.0) < STORAGE_WARN_INTERVAL: + return + self._storage_warned = now + _log(f"[cred] {label}持久化失败（{self._cooldowns.last_error}）；本次运行仍按内存态生效") + + def _forget_credential_cooldown(self, entry): + """Lift a persisted breaker while keeping this account's model cooldowns.""" + if self._durable_identity(entry): + outcome = self._cooldowns.clear_credential(entry["account_key"], entry["profile"]) + self._warn_storage("cooldown", outcome["durable"]) + + def cooldown_detail(self) -> list: + """Return persisted cooldown rows for diagnostics.""" + return self._cooldowns.detail() + + def cooldown_storage(self) -> dict: + """Report whether cooldown persistence is currently usable.""" + return {"available": self._cooldowns.path is not None, "path": self._cooldowns.path, + "degraded": self._cooldowns.last_error is not None, + "last_error": self._cooldowns.last_error, "rows": len(self._cooldowns.detail()), + "warning": "冷却持久化写入失败，本次运行仍按内存态生效。" if self._cooldowns.last_error else None} + + def clear_cooldowns(self, cm, model: str | None = None) -> dict: + """Lift a circuit breaker or model cooldown after a confirmed recovery or admin reset. + + Persisting cooldowns removes the old "restart the gateway to clear it" workaround, + so an explicit reset has to be able to lift one both in memory and on disk. The two + outcomes are reported separately: an in-memory reset that could not be written is + not a durable reset, and saying otherwise would hide a cooldown that comes back. + """ + with self._lock: + entry = next((e for e in self._entries if e["cm"] is cm), None) + if entry is None: + return {"changed_in_memory": False, "durable": False} + # The in-memory reset always happens, even when this account's identity is not + # complete enough to key durable state. + durable = not self._durable_identity(entry) + if model: + routed_model = _upstream_model(model, self._entry_profile(entry)) + changed = self._model_fail.pop((entry["id"], routed_model), None) is not None + if not durable: + outcome = self._cooldowns.clear_model(entry["account_key"], entry["profile"], routed_model) + durable = outcome["durable"] + self._warn_storage("cooldown", durable) + return {"changed_in_memory": changed, "durable": durable} + changed = entry["fail_until"] > time.time() + entry["fail_until"] = 0.0 + entry["last_error"] = None + if not durable: + durable = self._cooldowns.clear_credential(entry["account_key"], entry["profile"])["durable"] + self._warn_storage("cooldown", durable) + return {"changed_in_memory": changed, "durable": durable} + + def reset_cooldowns_for(self, identity: str) -> dict: + """Lift every cooldown held by one account, addressed by its public identity. + + Persisting cooldowns removed the old "restart the gateway to clear it" workaround, so an + operator needs a supported way back when a breaker or a 429 cooldown was recorded in + error, or when upstream has demonstrably recovered. The whole account is reset rather + than just its circuit breaker: a credential that is mid-429 is exactly the case an + operator is trying to unstick, and lifting only the breaker would leave it unusable. + + This is a local state change. It never refreshes a token, contacts upstream, or queues + synchronization, so a subsequent genuine failure is free to arm the cooldown again. + """ + with self._lock: + entry = next((e for e in self._entries if e.get("account_key") == identity), None) + if entry is None: + raise KeyError(identity) + now = time.time() + changed = entry["fail_until"] > now + entry["fail_until"] = 0.0 + entry["last_error"] = None + for key in [k for k in self._model_fail if k[0] == entry["id"]]: + if self._model_fail[key] > now: + changed = True + del self._model_fail[key] + # One durable write for the whole account, replacing any breaker and model rows. An + # account whose identity is incomplete still resets in memory, but owns no disk row. + durable = not self._durable_identity(entry) + if not durable: + durable = self._cooldowns.forget(entry["account_key"])["durable"] + self._warn_storage("cooldown", durable) + return {"changed_in_memory": changed, "durable": durable} + def entries(self) -> list[dict]: """Return credential snapshots for account maintenance.""" with self._lock: @@ -926,6 +1083,7 @@ def cooldown(self, cm: CredentialManager, reason: str = "", *, generation=None): e["fail_until"] = time.time() + CRED_COOLDOWN e["last_error"] = sanitize_log_text(reason, 256) e["last_failure_at"] = time.time() + self._remember_credential(e, e["last_error"]) _log(f"[cred] 凭证熔断 {CRED_COOLDOWN}s: {Path(cm.path).name} {reason}") def note_status(self, cm: CredentialManager | None, status: int, @@ -959,6 +1117,7 @@ def note_status(self, cm: CredentialManager | None, status: int, key = (e["id"], routed_model) until = max(until, self._model_fail.get(key, 0.0)) self._model_fail[key] = until + self._remember_model(e, routed_model, until) _log(f"[cred] 模型冷却 {model} @ {Path(cm.path).name} 至 " f"{time.strftime('%m-%d %H:%M:%S', time.localtime(until))} (HTTP 429)") @@ -1095,6 +1254,26 @@ def remove_file(self, name: str) -> bool: self.prune() return True + def forget_cooldowns(self, entry): + """Drop persisted cooldowns for a credential that no longer exists.""" + identity = entry.get("account_key") + if identity: + self._warn_storage("cooldown", self._cooldowns.forget(identity)["durable"]) + + def forget_usage(self, entry): + """Drop a deleted credential's cached usage, in the store and the live aggregate.""" + snapshots = CONFIG.get("usage_snapshots") + if snapshots is not None: + snapshots.forget(entry["id"]) + accounts = CONFIG.get("usage_daily_accounts") + if isinstance(accounts, dict): + accounts.pop(entry["id"], None) + + def forget_credential_state(self, entry): + """Run both independent cleanups for a credential that is gone or replaced.""" + self.forget_cooldowns(entry) + self.forget_usage(entry) + def first(self) -> CredentialManager | None: with self._lock: return self._entries[0]["cm"] if self._entries else None @@ -1336,11 +1515,20 @@ def _sync_usage(pool, entries=None, expected_identity=None): usage = credits_mod.fetch_request_usage(_bearer_token(headers), uid=headers.get("X-User-Id", ""), domain=headers.get("X-Domain", "")) def store(): - accounts[entry["id"]] = {"site": site, "by_day": usage["by_day"], + accounts[entry["id"]] = {"identity": entry.get("account_key"), "site": site, + "by_day": usage["by_day"], "total_credits": round(usage["total_credits"], 2), "requests": usage["requests"], "partial": bool(usage.get("partial")), "fetched_at": time.time()} + snapshots = CONFIG.get("usage_snapshots") + if snapshots is not None and entry.get("account_key") and entry.get("uid"): + # Only a validated identity may key durable state. An absent flag defaults to + # False, but a present value is passed through raw so the store's own strict + # validator governs: coercing it here with bool() would mask a malformed value + # and persist it as a legitimate flag. + snapshots.store(entry["id"], entry["account_key"], site, usage, + partial=usage.get("partial", False)) if not pool.apply_if_current(cm, generation, store): stale.add(entry["id"]) except Exception as error: @@ -1350,48 +1538,114 @@ def store(): return stale & target_ids +def _usage_row_expired(snap) -> bool: + """Whether a cached row is too old to keep showing.""" + return time.time() - float(snap.get("fetched_at") or 0.0) > USAGE_CACHE_MAX_AGE_S + + +def _adopt_cached_usage(pool, snapshots, accounts): + """Seed the aggregate from the cache, keeping only rows that still belong to their account. + + Ownership is checked here, at the moment of use, rather than only when the cache was + written: a path can be reused between restarts, and a row hydrated on an earlier pass + would otherwise keep showing the previous account's usage indefinitely. + """ + # The live map is mutated by forget_usage under the pool lock, so the reads below need that + # same lock to stay consistent with a concurrent credential deletion or replacement. + with pool._lock: + cached = snapshots.accounts() + if not cached and not accounts: + return # Nothing cached and nothing live: leave the pool untouched. + # Only an account whose identity components are all present may own durable usage: an + # account_key is a hash of the profile and UID, so an account with no UID would otherwise + # share one identity with every other such account. + owners = {entry["id"]: entry.get("account_key") for entry in pool.entries() + if entry.get("account_key") and entry.get("profile") and entry.get("uid")} + # Drop any row whose recorded owner no longer matches the account at that path. This covers + # rows hydrated by an earlier pass as well as live rows, so a reused path cannot keep + # displaying the previous account's usage. Iterate over a copy of the items, so a row is + # never looked up in the live map after the keys were copied. + for path, row in list(accounts.items()): + if owners.get(path) != row.get("identity"): + accounts.pop(path, None) + for path, row in cached.items(): + identity = row.get("identity") + if owners.get(path) != identity: + snapshots.forget(path) # The path moved on; the cached row is not ours to show. + continue + if _usage_row_expired(row): + snapshots.forget(path) + continue + # A live snapshot from this run always wins; the cache only fills a gap. + if path in accounts: + continue + # A restored snapshot is stale until a refresh confirms it; clearing that is per account. + accounts[path] = dict(row) + accounts[path]["stale"] = True + CONFIG["usage_daily_accounts"] = accounts + + def _publish_usage_daily(pool, stale=()): """Aggregate enabled accounts' usage, retaining failed snapshots with explicit staleness.""" - accounts = CONFIG.get("usage_daily_accounts") - if not isinstance(accounts, dict): - accounts = {} - enabled = {e["id"] for e in pool.entries() if model_policy.credential_enabled(CONFIG, e)} - by_day, groups = {}, {} - used, count = 0.0, 0 - partial = False - newest = 0.0 - stale_out = [] - for cred_id, snap in accounts.items(): - if cred_id not in enabled: - continue - site = snap.get("site") or "domestic" - group = groups.setdefault(site, {"by_day": {}, "total_credits": 0.0, "requests": 0}) - for day, models in (snap.get("by_day") or {}).items(): - total_day = by_day.setdefault(day, {}) - site_day = group["by_day"].setdefault(day, {}) - for model, credit in models.items(): - total_day[model] = round(total_day.get(model, 0.0) + credit, 6) - site_day[model] = round(site_day.get(model, 0.0) + credit, 6) - group["total_credits"] += float(snap.get("total_credits") or 0) - group["requests"] += int(snap.get("requests") or 0) - used += float(snap.get("total_credits") or 0) - count += int(snap.get("requests") or 0) - newest = max(newest, float(snap.get("fetched_at") or 0)) - if snap.get("partial"): - partial = True - # Failed enabled accounts must remain visible even without a prior snapshot. - for cred_id in stale: - if cred_id in enabled: - partial = True - stale_out.append(Path(cred_id).name) - # A zero timestamp preserves quota-difference fallback when no usage snapshot exists. - for group in groups.values(): - group["total_credits"] = round(group["total_credits"], 2) - out = {"by_day": by_day, "groups": groups, "total_credits": round(used, 2), - "requests": count, "fetched_at": newest, "partial": partial} - if stale_out: - out["stale_accounts"] = sorted(stale_out) - CONFIG["usage_daily"] = out + # CredentialPool.forget_usage removes entries from the live map under the pool lock, so the + # snapshot taken here and the pruning inside the loop must hold that same lock. Copying the + # keys and then indexing the live dictionary would otherwise raise KeyError when a credential + # is deleted or replaced while a usage-maintenance pass is running. + with pool._lock: + accounts = CONFIG.get("usage_daily_accounts") + if not isinstance(accounts, dict): + accounts = {} + CONFIG["usage_daily_accounts"] = accounts + # Seed from the on-disk cache so the dashboard is not blank until the first refresh. + snapshots = CONFIG.get("usage_snapshots") + if snapshots is not None: + _adopt_cached_usage(pool, snapshots, accounts) + enabled = {e["id"] for e in pool.entries() if model_policy.credential_enabled(CONFIG, e)} + # Aggregate from a consistent local copy of the live map rather than indexing it per row. + rows = dict(accounts) + by_day, groups = {}, {} + used, count = 0.0, 0 + partial = False # Upstream paging hid additional usage for some account. + stale_out = set() # Accounts whose displayed figures are not from this run. + newest = 0.0 + for cred_id, snap in rows.items(): + if cred_id not in enabled: + continue + if _usage_row_expired(snap): + accounts.pop(cred_id, None) # Drop it from the live map, not just from this sum. + continue + site = snap.get("site") or "domestic" + group = groups.setdefault(site, {"by_day": {}, "total_credits": 0.0, "requests": 0}) + for day, models in (snap.get("by_day") or {}).items(): + total_day = by_day.setdefault(day, {}) + site_day = group["by_day"].setdefault(day, {}) + for model, credit in models.items(): + total_day[model] = round(total_day.get(model, 0.0) + credit, 6) + site_day[model] = round(site_day.get(model, 0.0) + credit, 6) + group["total_credits"] += float(snap.get("total_credits") or 0) + group["requests"] += int(snap.get("requests") or 0) + used += float(snap.get("total_credits") or 0) + count += int(snap.get("requests") or 0) + newest = max(newest, float(snap.get("fetched_at") or 0)) + # `partial` is the aggregate "this view is incomplete" flag the dashboard shows, so a + # failed or not-yet-refreshed account sets it; `stale_accounts` names which ones. + if snap.get("partial") or snap.get("stale"): + partial = True + if snap.get("stale"): + stale_out.add(Path(cred_id).name) + # Failed enabled accounts must remain visible even without a prior snapshot. + for cred_id in stale: + if cred_id in enabled: + partial = True + stale_out.add(Path(cred_id).name) + # A zero timestamp preserves quota-difference fallback when no usage snapshot exists. + for group in groups.values(): + group["total_credits"] = round(group["total_credits"], 2) + out = {"by_day": by_day, "groups": groups, "total_credits": round(used, 2), + "requests": count, "fetched_at": newest, "partial": partial} + if stale_out: + out["stale_accounts"] = sorted(stale_out) + CONFIG["usage_daily"] = out _log(f"[usage] 明细已同步: {count} 请求 / {used:.2f} credits" + (f" | {len(stale_out)} 账号同步失败" if stale_out else "")) @@ -1402,6 +1656,12 @@ def _housekeep_once(pool: CredentialPool, ledger, *, pending_only=False): return with _HOUSEKEEP_LOCK: pool._rescan() + # Drop expired deadlines and aged usage so neither table grows without bound; a failed + # prune is reported rather than silently leaving stale rows on disk. + pool._cooldowns.prune() + pool._warn_storage("cooldown", pool._cooldowns.last_error is None) + if CONFIG.get("usage_snapshots") is not None: + CONFIG["usage_snapshots"].prune() ids = pool.begin_sync(all_entries=not pending_only) failed = set() try: @@ -1525,6 +1785,7 @@ async def _protocol_http_exception(request: Request, exc: HTTPException): "retry_write_timeout": False, # Opt-in replay after incomplete writes "usage_daily": None, # Usage aggregated by date and model "usage_daily_accounts": None, # Independent per-account usage snapshots + "usage_snapshots": None, # On-disk cache of the per-account snapshots "credit_price_cny": None, "credit_price_usd": None, "usd_rate": None, "desensitize": False, "no_compact": False, "keep_tool_metadata": False} # None prices use module defaults. @@ -1933,11 +2194,14 @@ async def admin_credential_action(identity: str, action: str, request: Request, from app.admin_api import _body try: body = await _body(request, 4096, allow_empty=True) - if body and (action != "travel" or set(body) != {"confirm_buddy", "agreement_revision"} - or body["confirm_buddy"] is not True or not isinstance(body["agreement_revision"], str)): - raise ValueError() except ValueError: - raise HTTPException(400, "首领确认参数无效") from None + raise HTTPException(400, "请求体必须是有效 JSON 对象") from None + if body: + if action != "travel": + raise HTTPException(400, "该操作不接受请求体") + if (set(body) != {"confirm_buddy", "agreement_revision"} + or body["confirm_buddy"] is not True or not isinstance(body["agreement_revision"], str)): + raise HTTPException(400, "首领确认参数无效") return await run_in_threadpool(_admin_credential_action, action, identity, consent_revision=body.get("agreement_revision")) @@ -3637,7 +3901,9 @@ def main(): if not files: seed_credentials() # Seed missing desktop credentials into managed storage. CONFIG["cred_pool"] = CredentialPool(files, scan=not files, - blocks_path=managed_auth_dir() / "model-site-blocks.json") + blocks_path=managed_auth_dir() / "model-site-blocks.json", + cooldowns_path=managed_auth_dir() / "credential-cooldowns.json") + CONFIG["usage_snapshots"] = UsageSnapshots(managed_auth_dir() / "usage-snapshots.json") CONFIG["cred"] = CONFIG["cred_pool"].first() CONFIG["account_catalogs"] = {} # Disable static fallback before maintenance starts. if credits_mod is not None: @@ -3647,6 +3913,12 @@ def main(): managed_auth_dir() / "model-catalog.json", ttl=args.model_catalog_ttl) CONFIG["cred_pool"].set_ledger(ledger) # Verify balance ownership before publishing catalogs. _publish_model_cache() + # Publish cached usage before maintenance threads start, so the dashboard is populated from + # the first request. It stays a no-op when nothing is cached, so an empty deployment and a + # mocked pool are both unaffected. + if CONFIG["usage_snapshots"].detail(): + _publish_usage_daily(CONFIG["cred_pool"]) + runtime_management.install(sys.modules[__name__]) threading.Thread(target=_refresher_loop, args=(CONFIG["cred_pool"],), daemon=True, name="cred-refresher").start() @@ -3671,6 +3943,7 @@ def main(): sys.stderr.write(" GET /admin/credits (积分/签到状态)\n") sys.stderr.write(" POST /admin/checkin (仅签到，按日幂等)\n") sys.stderr.write(" POST /admin/sync (同步余额、目录与用量，不签到)\n") + sys.stderr.write(" POST /admin/credentials/{id}/reset-cooldown (清除该账号冷却，仅本地)\n") sys.stderr.write(" 每日签到 + 快过期积分优先调度已启用\n") if args.api_key: sys.stderr.write(" 鉴权已启用（API key 已设置）\n") diff --git a/docs/advanced.md b/docs/advanced.md index 0fa3928..f1655e0 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -110,10 +110,12 @@ Trial credits are manual-only for eligible `intl-work` accounts: use the credent | `POST /admin/oauth/start` · `GET /admin/oauth/poll` | Start/poll login; `site=cn` (default), `intl` (international WorkBuddy) or `intl-codebuddy` (international CodeBuddy) | | `GET /admin/credits` · `POST /admin/checkin` | Inspect credits; daily-idempotent check-in followed by domestic travel when enabled | | `POST /admin/sync` | Synchronize all enabled accounts' balances, catalogs and usage; no check-in or trial claims | -| `POST /admin/credentials/{id}/{action}` | Single-account `refresh`, `checkin`, `sync`, `travel-status` (query only), `travel` (claim then dispatch), or `trial` (one-time trial credits) | +| `POST /admin/credentials/{id}/{action}` | Single-account `refresh`, `checkin`, `sync`, `travel-status` (query only), `travel` (claim then dispatch), `trial` (one-time trial credits), or `reset-cooldown` (local-only) | Travel results include `phase`, optional safe `error_kind`/`http_status`/`code`, and snapshot `remaining_seconds`. `claimed`/`departed` remain true for confirmed writes even if a later query sets `ok=false` and `stale=true`; query status before another attempt. +`reset-cooldown` takes no request body and lifts every cooldown held by one account — both its 401/403 circuit breaker and its per-model 429 cooldowns — so an operator can recover from a cooldown recorded in error without restarting the gateway. It only edits local state: it never refreshes a token, contacts upstream, or queues synchronization, and it works for a manually disabled account. The result separates `changed_in_memory` from `durable`; `ok` is false when the write failed, in which case the in-memory reset is already effective but a restart would restore the stored row, and the request can simply be repeated. + Pages use `/dashboard/*`, management APIs use `/admin/*`, and clients retain `/v1/*`. `/cn` and `/intl` API prefixes are not registered. Automatic model routing requires no client URL changes. Management requires an API key. The WebUI exchanges that key for an HttpOnly management Cookie, which only authorizes `/admin/*`, not `/v1/*`. API clients send `Authorization: Bearer ` or `X-Api-Key`. An empty key preserves legacy unauthenticated inference only, not management. `/health` never exposes account, path or exception details. diff --git a/docs/advanced.zh-CN.md b/docs/advanced.zh-CN.md index 9f111b1..4c9cc6b 100644 --- a/docs/advanced.zh-CN.md +++ b/docs/advanced.zh-CN.md @@ -110,10 +110,12 @@ scoped 模式可选传入 `X-Codebuddy-Session-ID`、`metadata.conversation_id` | `POST /admin/oauth/start` · `GET /admin/oauth/poll` | 发起与轮询登录；`site=cn`（默认）、`intl`（国际 WorkBuddy）或 `intl-codebuddy`（国际 CodeBuddy） | | `GET /admin/credits` · `POST /admin/checkin` | 查询额度；按日幂等签到，国内按开关继续旅行 | | `POST /admin/sync` | 同步全部启用账号的余额、目录和用量，不签到、不领取试用 | -| `POST /admin/credentials/{id}/{action}` | 单账号 `refresh`、`checkin`、`sync`、`travel-status`（仅查询）、`travel`（领取后派出）或 `trial`（一次性体验积分） | +| `POST /admin/credentials/{id}/{action}` | 单账号 `refresh`、`checkin`、`sync`、`travel-status`（仅查询）、`travel`（领取后派出）、`trial`（一次性体验积分）或 `reset-cooldown`（仅本地状态） | 旅行结果包含 `phase`、可选的安全诊断 `error_kind`/`http_status`/`code` 和查询时的 `remaining_seconds`。后续查询失败会设置 `ok=false`、`stale=true`，但保留已确认的 `claimed`/`departed`；再次操作前先查询核验。 +`reset-cooldown` 不接受请求体，会清除该账号的全部冷却：既包括 401/403 认证熔断，也包括按模型的 429 冷却。这样在冷却被误判或上游已恢复时，无需重启网关即可放行。它只修改本地状态：不刷新 Token、不访问上游、不排队同步，人工停用的账号同样可用。结果区分 `changed_in_memory` 与 `durable`；写入失败时 `ok` 为 false，此时内存冷却已清除但重启会恢复盘上记录，直接重试该请求即可。 + 页面使用 `/dashboard/*`，管理 API 使用 `/admin/*`，客户端保留原 `/v1/*`；不注册 `/cn`、`/intl` API 前缀。模型自动选路不要求客户端改变地址。 管理必须配置 API key；WebUI 使用同 key 建立 HttpOnly 管理 Cookie，Cookie 仅授权 `/admin/*`，不能用于 `/v1/*`。命令行 API 请求携带 `Authorization: Bearer ` 或 `X-Api-Key`。空 key 仅保留推理接口的历史无鉴权行为，不开放管理；`/health` 不返回账号、路径或异常详情。 diff --git a/tests/test_credential_actions.py b/tests/test_credential_actions.py index eaec48a..486b80a 100644 --- a/tests/test_credential_actions.py +++ b/tests/test_credential_actions.py @@ -3,6 +3,7 @@ from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +import json import time import unittest from unittest.mock import patch @@ -179,5 +180,214 @@ def test_api_and_cookie_auth_require_csrf_for_new_actions(self): self.assertTrue(self.post("checkin")["ok"]) +class ResetCooldownTests(unittest.TestCase): + """The admin reset is the supported replacement for restarting the gateway to clear one.""" + + add_account = fixtures.RegionRoutingTests.add_account + configure = fixtures.RegionRoutingTests.configure + handle_upstream = fixtures.RegionRoutingTests.handle_upstream + + def setUp(self): + fixtures.RegionRoutingTests.setUp(self) + converter.CONFIG.update(api_key="synthetic-management-key", usage_daily_accounts={}, + usage_daily={}, cooldowns_path=self.root / "cooldowns.json") + self.control = ControlStore(self.root / "control.sqlite3") + self.audit = AuditStore(self.root / "audit.sqlite3") + self.addCleanup(self.control.close) + self.addCleanup(self.audit.close) + converter.CONFIG.update(control_store=self.control, audit_store=self.audit) + app = FastAPI() + app.router.routes = list(converter.app.router.routes) + install_admin(app, converter.CONFIG, Management(converter)) + self.client = self.enterContext(TestClient(app, base_url="https://testserver", + headers={"Authorization": "Bearer synthetic-management-key"})) + + def build_pool(self, *, cooldowns_path=None): + """Build a fresh pool bound to a real cooldown file, as a restart would.""" + path = cooldowns_path if cooldowns_path is not None else self.root / "cooldowns.json" + pool = converter.CredentialPool( + [self.root / (profile + ".info") for profile in ("cn-cli", "intl-work")], + cooldowns_path=path) + converter.CONFIG["cred_pool"] = pool + return pool + + def arm(self, pool): + """Record both kinds of cooldown on the first account.""" + entry = pool._entries[0] + pool.cooldown(entry["cm"], reason="backend HTTP 401") + pool.note_status(entry["cm"], 429, model="glm-5.3-flash", raw=b"") + return entry + + def url(self, entry): + return "/admin/credentials/" + entry["account_key"] + "/reset-cooldown" + + def test_reset_lifts_both_cooldown_kinds_and_survives_a_restart(self): + pool = self.build_pool() + entry = self.arm(pool) + self.assertFalse(pool._healthy(entry)) + self.assertFalse(pool._model_healthy(entry, "glm-5.3-flash")) + + response = self.client.post(self.url(entry)) + self.assertEqual(response.status_code, 200, response.text) + result = response.json()["results"][0] + self.assertTrue(result["ok"], result) + self.assertTrue(result["changed_in_memory"]) + self.assertTrue(result["durable"]) + # The same live pool must be usable immediately; checking only a rebuilt pool would miss + # an in-memory model cooldown that was never lifted. + self.assertTrue(pool._healthy(entry)) + self.assertTrue(pool._model_healthy(entry, "glm-5.3-flash")) + + # A brand new pool reads the file back, which is what a restart does. + revived = self.build_pool()._entries[0] + self.assertTrue(self.build_pool()._healthy(revived)) + self.assertTrue(self.build_pool()._model_healthy(revived, "glm-5.3-flash")) + self.assertIsNone(revived.get("last_error")) + + def test_reset_never_contacts_upstream_or_refreshes_tokens(self): + pool = self.build_pool() + entry = self.arm(pool) + before = entry["cm"]._generation + with patch.object(converter.credits_mod, "fetch_credits") as credits_call, \ + patch.object(converter, "_sync_credits") as sync: + self.client.post(self.url(entry)) + self.assertEqual(self.requests, []) # No upstream call at all. + credits_call.assert_not_called() + sync.assert_not_called() + self.assertEqual(entry["cm"]._generation, before) # No token refresh. + + def test_reset_is_allowed_while_maintenance_holds_the_lock(self): + """A local state edit must not queue behind the hourly sweep.""" + pool = self.build_pool() + entry = self.arm(pool) + with converter._HOUSEKEEP_LOCK: + response = self.client.post(self.url(entry)) + self.assertEqual(response.status_code, 200, response.text) + self.assertTrue(response.json()["results"][0]["ok"]) + + def test_reset_works_without_a_ledger(self): + pool = self.build_pool() + entry = self.arm(pool) + converter.CONFIG["ledger"] = None + response = self.client.post(self.url(entry)) + self.assertEqual(response.status_code, 200, response.text) + self.assertTrue(response.json()["results"][0]["ok"]) + + def test_reset_works_for_a_manually_disabled_account(self): + pool = self.build_pool() + entry = self.arm(pool) + Management(converter).admin_set_credential_enabled(entry["account_key"], False) + response = self.client.post(self.url(entry)) + self.assertEqual(response.status_code, 200, response.text) + result = response.json()["results"][0] + self.assertTrue(result["ok"], result) + self.assertNotIn("skipped", result) + # _healthy() also requires manual enablement, so check the cooldown itself: the reset + # must clear it even though the account stays disabled. + self.assertEqual(self.build_pool()._entries[0]["fail_until"], 0.0) + + def test_reset_leaves_other_accounts_untouched(self): + pool = self.build_pool() + first = self.arm(pool) + other = pool._entries[1] + pool.cooldown(other["cm"], reason="backend HTTP 403") + self.client.post(self.url(first)) + revived = self.build_pool() + others = {e["account_key"]: e for e in revived._entries} + self.assertGreater(others[other["account_key"]]["fail_until"], time.time()) + + def test_reset_of_an_unknown_identity_is_rejected(self): + self.build_pool() + response = self.client.post("/admin/credentials/" + "0" * 64 + "/reset-cooldown") + self.assertEqual(response.status_code, 404) + self.assertEqual(self.requests, []) + + def test_reset_rejects_a_replaced_identity_instead_of_clearing_the_new_account(self): + pool = self.build_pool() + entry = self.arm(pool) + stale = entry["account_key"] + # The file now holds a different account, so the old identity no longer resolves. + self.add_account("replacement-uid", "cn-cli") + (self.root / "cn-cli.info").write_text(json.dumps(self.credentials["replacement-uid"]), + encoding="utf-8") + response = self.client.post("/admin/credentials/" + stale + "/reset-cooldown") + self.assertEqual(response.status_code, 404) + self.assertEqual(self.requests, []) + + def test_a_failed_write_reports_failure_and_stays_retryable(self): + pool = self.build_pool() + entry = self.arm(pool) + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + response = self.client.post(self.url(entry)) + self.assertEqual(response.status_code, 200, response.text) + result = response.json()["results"][0] + self.assertFalse(result["ok"]) # Never report a failed write as done. + self.assertTrue(result["changed_in_memory"]) + self.assertFalse(result["durable"]) + self.assertIn("写入失败", result["message"]) + self.assertTrue(pool._healthy(entry)) # In-memory effect is real. + # The disk row is still there, so a restart restores it... + self.assertFalse(self.build_pool()._healthy(self.build_pool()._entries[0])) + # ...and the retry (with the fault removed) clears it for good. + retry = self.client.post(self.url(entry)).json()["results"][0] + self.assertTrue(retry["ok"], retry) + self.assertTrue(self.build_pool()._healthy(self.build_pool()._entries[0])) + + def test_reset_requires_authentication(self): + pool = self.build_pool() + entry = self.arm(pool) + self.client.headers.pop("authorization") + self.assertEqual(self.client.post(self.url(entry)).status_code, 401) + self.assertFalse(pool._healthy(entry)) # Rejected without mutating. + + def test_reset_requires_csrf_and_same_origin_for_cookie_sessions(self): + pool = self.build_pool() + entry = self.arm(pool) + self.client.headers.pop("authorization") + session = self.client.post("/admin/session", json={"api_key": "synthetic-management-key"}, + headers={"Origin": "https://testserver"}) + self.assertEqual(session.status_code, 200, session.text) + self.assertEqual(self.client.post(self.url(entry)).status_code, 403) # No CSRF yet. + self.assertFalse(pool._healthy(entry)) + self.client.headers["X-CSRF-Token"] = session.json()["csrf_token"] + self.client.headers["Origin"] = "https://evil.example" + self.assertEqual(self.client.post(self.url(entry)).status_code, 403) # Cross-site. + self.assertFalse(pool._healthy(entry)) + self.client.headers["Origin"] = "https://testserver" + self.assertEqual(self.client.post(self.url(entry)).status_code, 200) + self.assertTrue(pool._healthy(entry)) + + def test_reset_rejects_a_request_body(self): + pool = self.build_pool() + entry = self.arm(pool) + response = self.client.post(self.url(entry), json={"model": "glm-5.3-flash"}) + self.assertEqual(response.status_code, 400) + self.assertFalse(pool._healthy(entry)) # Rejected without mutating. + + def test_unknown_action_is_still_rejected(self): + self.build_pool() + self.assertEqual(self.client.post("/admin/credentials/missing/reset-cooldowns").status_code, 404) + + def test_reset_with_nothing_recorded_is_a_successful_no_op(self): + pool = self.build_pool() + entry = pool._entries[0] + response = self.client.post(self.url(entry)) + self.assertEqual(response.status_code, 200, response.text) + result = response.json()["results"][0] + self.assertTrue(result["ok"], result) + self.assertFalse(result["changed_in_memory"]) + self.assertTrue(result["durable"]) + + def test_reset_is_recorded_in_the_audit_trail(self): + pool = self.build_pool() + entry = self.arm(pool) + self.client.post(self.url(entry)) + events = self.audit.list_records("admin", limit=50)["items"] + matching = [e for e in events if e.get("action") == "credential.reset-cooldown"] + self.assertTrue(matching, events) + self.assertEqual(matching[0]["details"]["stage"], "durable") + self.assertEqual(matching[0]["details"]["outcome"], "success") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_credential_cooldowns.py b/tests/test_credential_cooldowns.py new file mode 100644 index 0000000..6c5e4f7 --- /dev/null +++ b/tests/test_credential_cooldowns.py @@ -0,0 +1,524 @@ +#!/usr/bin/env python3 +"""Test that credential cooldowns survive a restart without ever over-extending them.""" + +import json +import os +import sys +import tempfile +import time +import unittest +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +import converter +from app.credential_cooldowns import (AUTH_CEILING_S, MAX_ACCOUNTS, MAX_BYTES, MAX_MODELS, + MODEL_CEILING_S, VERSION, CredentialCooldowns) + +IDENTITY = "a" * 64 +OTHER_IDENTITY = "b" * 64 +PROFILE = "cn-cli" +OTHER_PROFILE = "intl-cli" +MODEL = "claude-sonnet-4" + + +class CooldownStoreTests(unittest.TestCase): + """Exercise the store directly, with no credential files and no network access.""" + + def setUp(self): + self.root = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.path = self.root / "credential-cooldowns.json" + + def row(self, **overrides): + row = {"profile": PROFILE, "fail_until": time.time() + 300, "reason": "backend HTTP 401", + "failed_at": time.time(), "models": {}} + row.update(overrides) + return row + + def write(self, accounts): + self.path.write_text(json.dumps({"version": VERSION, "accounts": accounts}), encoding="utf-8") + + def test_credential_cooldown_survives_a_restart(self): + first = CredentialCooldowns(self.path) + self.assertTrue(first.note_credential(IDENTITY, PROFILE, time.time() + 300, reason="backend HTTP 401")) + restored = CredentialCooldowns(self.path) + self.assertGreater(restored.credential_until(IDENTITY, PROFILE), time.time()) + self.assertEqual(restored.restore(IDENTITY, PROFILE)["reason"], "backend HTTP 401") + + def test_model_cooldown_survives_a_restart_and_stays_per_model(self): + first = CredentialCooldowns(self.path) + first.note_model(IDENTITY, PROFILE, MODEL, time.time() + 600) + restored = CredentialCooldowns(self.path) + self.assertGreater(restored.model_until(IDENTITY, PROFILE, MODEL), time.time()) + self.assertEqual(restored.model_until(IDENTITY, PROFILE, "another-model"), 0.0) + + def test_repeated_restarts_never_extend_a_deadline(self): + first = CredentialCooldowns(self.path) + first.note_credential(IDENTITY, PROFILE, time.time() + 300) + original = json.loads(self.path.read_text())["accounts"][IDENTITY]["fail_until"] + for _ in range(5): + store = CredentialCooldowns(self.path) + store.note_credential(IDENTITY, PROFILE, store.credential_until(IDENTITY, PROFILE)) + self.assertLessEqual(json.loads(self.path.read_text())["accounts"][IDENTITY]["fail_until"], + original + 0.01) + + def test_the_writer_never_produces_a_file_the_reader_rejects(self): + """A file the reader discards would lose every cooldown it held, not just one row.""" + now = time.time() + store = CredentialCooldowns(self.path) + # Fill the table to capacity in memory, then write once: the byte bound must hold + # even at the worst case of MAX_ACCOUNTS x MAX_MODELS x longest strings. + with store._lock: + for index in range(MAX_ACCOUNTS): + identity = f"{index:064x}" + store._data[identity] = {"profile": PROFILE, "fail_until": now + 300, + "reason": "x" * 256, "failed_at": now, + "models": {f"model-name-{model:03d}": now + 600 + for model in range(MAX_MODELS)}} + content, shed = store._serialize_locked() + self.assertLessEqual(len(content), MAX_BYTES) + self.assertTrue(shed) # Capacity really did force shedding. + self.path.write_bytes(content) + restored = CredentialCooldowns(self.path) + self.assertGreater(len(restored.detail()), 0) # The reader still accepts it. + self.assertLessEqual(self.path.stat().st_size, MAX_BYTES) + + def test_auth_and_model_ceilings_are_enforced_separately(self): + store = CredentialCooldowns(self.path) + now = time.time() + store.note_credential(IDENTITY, PROFILE, now + MODEL_CEILING_S) # A quota-length deadline. + store.note_model(IDENTITY, PROFILE, MODEL, now + 100 * MODEL_CEILING_S) + self.assertLessEqual(store.credential_until(IDENTITY, PROFILE) - now, AUTH_CEILING_S + 1) + self.assertLessEqual(store.model_until(IDENTITY, PROFILE, MODEL) - now, MODEL_CEILING_S + 1) + + def test_a_deadline_beyond_the_ceiling_is_never_adopted(self): + now = time.time() + cases = {"far future": 1e12, "past the auth ceiling": now + AUTH_CEILING_S + 3600, + "in the past": now - 60, "bool": True, "string": "nan", "infinity": float("inf")} + for label, until in cases.items(): + with self.subTest(fail_until=label): + self.write({IDENTITY: self.row(fail_until=until)}) + self.assertEqual(CredentialCooldowns(self.path).credential_until(IDENTITY, PROFILE), 0.0) + for label, until in {"far future": 1e12, "past the model ceiling": now + MODEL_CEILING_S + 3600, + "huge integer": 10 ** 400}.items(): + with self.subTest(models=label): + self.write({IDENTITY: self.row(models={MODEL: until})}) + self.assertEqual(CredentialCooldowns(self.path).model_until(IDENTITY, PROFILE, MODEL), 0.0) + + def test_expired_rows_are_dropped_rather_than_restored(self): + self.write({IDENTITY: self.row(fail_until=time.time() - 1, models={MODEL: time.time() - 1})}) + store = CredentialCooldowns(self.path) + self.assertEqual(store.detail(), []) + self.assertEqual(store.restore(IDENTITY, PROFILE), {}) + + def test_prune_removes_dead_rows_and_keeps_live_ones(self): + store = CredentialCooldowns(self.path) + store.note_model(IDENTITY, PROFILE, "live-model", time.time() + 600) + with store._lock: # Write an already-expired breaker directly. + store._data[IDENTITY]["fail_until"] = time.time() - 1 + self.assertTrue(store.prune()["changed"]) + self.assertEqual(set(store.restore(IDENTITY, PROFILE)["models"]), {"live-model"}) + self.assertEqual(store.credential_until(IDENTITY, PROFILE), 0.0) + + def test_a_full_table_reports_shedding_instead_of_claiming_success(self): + """A caller must not be told its cooldown is durable when it was shed for capacity.""" + now = time.time() + store = CredentialCooldowns(self.path) + with store._lock: + for index in range(MAX_ACCOUNTS): + store._data[f"{index:064x}"] = {"profile": PROFILE, "fail_until": now + 300, + "reason": "x" * 256, "failed_at": now, + "models": {f"model-name-{model:03d}": now + 600 + for model in range(MAX_MODELS)}} + self.assertFalse(store.note_model("f" * 64, PROFILE, "new-model", now + 600)) + self.assertEqual(store.last_error, "capacity") + + def test_invalid_deadlines_are_rejected_without_raising(self): + store = CredentialCooldowns(self.path) + for label, until in (("huge integer", 10 ** 400), ("nan", float("nan")), ("none", None), + ("string", "abc"), ("infinity", float("inf"))): + with self.subTest(until=label): + self.assertFalse(store.note_credential(IDENTITY, PROFILE, until)) + self.assertFalse(store.note_model(IDENTITY, PROFILE, MODEL, until)) + self.assertEqual(store.detail(), []) + + def test_untrusted_snapshots_are_rejected_wholesale(self): + payloads = { + "not json": b"nope", + "empty object": b"{}", + "unsupported version": json.dumps({"version": 99, "accounts": {}}).encode(), + "boolean version": json.dumps({"version": True, "accounts": {}}).encode(), + "extra top-level field": json.dumps({"version": VERSION, "accounts": {}, "x": 1}).encode(), + "non-hex identity": json.dumps({"version": VERSION, "accounts": {"zzz": self.row()}}).encode(), + "extra row field": json.dumps({"version": VERSION, + "accounts": {IDENTITY: {**self.row(), "extra": 1}}}).encode(), + "duplicate keys": b'{"version":1,"version":1,"accounts":{}}', + "infinity literal": ('{"version":1,"accounts":{"' + IDENTITY + '":{"profile":"cn-cli",' + '"fail_until":Infinity,"reason":"","failed_at":0,"models":{}}}}').encode(), + "bad profile token": json.dumps({"version": VERSION, "accounts": { + IDENTITY: {**self.row(), "profile": "bad profile!"}}}).encode(), + "oversized file": b"x" * (MAX_BYTES + 1), + } + for label, content in payloads.items(): + with self.subTest(payload=label): + self.path.write_bytes(content) + self.assertEqual(CredentialCooldowns(self.path).detail(), []) + + def test_oversized_tables_are_rejected_on_load(self): + accounts = {f"{index:064x}": self.row() for index in range(MAX_ACCOUNTS + 1)} + self.write(accounts) + self.assertEqual(CredentialCooldowns(self.path).detail(), []) + self.write({IDENTITY: self.row(models={f"m{index}": time.time() + 300 + for index in range(MAX_MODELS + 1)})}) + self.assertEqual(CredentialCooldowns(self.path).detail(), []) + + def test_writes_reject_values_the_reader_would_refuse(self): + store = CredentialCooldowns(self.path) + now = time.time() + self.assertFalse(store.note_credential("not-a-hash", PROFILE, now + 300)) + self.assertFalse(store.note_credential(IDENTITY, "bad profile!", now + 300)) + self.assertFalse(store.note_model(IDENTITY, PROFILE, "bad\nmodel", now + 300)) + self.assertEqual(store.detail(), []) + + def test_slash_qualified_model_names_round_trip(self): + """Vendor-qualified identifiers may contain slashes; the reader must accept them.""" + store = CredentialCooldowns(self.path) + store.note_model(IDENTITY, PROFILE, "vendor/model-v2", time.time() + 600) + self.assertGreater(CredentialCooldowns(self.path).model_until(IDENTITY, PROFILE, "vendor/model-v2"), + time.time()) + + def test_reused_path_does_not_inherit_another_products_cooldowns(self): + store = CredentialCooldowns(self.path) + store.note_credential(IDENTITY, PROFILE, time.time() + 300) + store.note_credential(IDENTITY, OTHER_PROFILE, time.time() + 300) + self.assertEqual(store.credential_until(IDENTITY, PROFILE), 0.0) + self.assertGreater(store.credential_until(IDENTITY, OTHER_PROFILE), time.time()) + + def test_clear_and_forget_remove_state(self): + store = CredentialCooldowns(self.path) + store.note_credential(IDENTITY, PROFILE, time.time() + 300) + store.note_model(IDENTITY, PROFILE, MODEL, time.time() + 600) + self.assertTrue(store.clear_model(IDENTITY, PROFILE, MODEL)["changed"]) + self.assertEqual(store.restore(IDENTITY, PROFILE)["models"], {}) + self.assertTrue(store.clear_credential(IDENTITY, PROFILE)["changed"]) + self.assertEqual(store.detail(), []) + store.note_credential(IDENTITY, PROFILE, time.time() + 300) + self.assertTrue(store.forget(IDENTITY)["changed"]) + self.assertEqual(json.loads(self.path.read_text())["accounts"], {}) + + def test_a_write_failure_is_reported_and_keeps_state_in_memory(self): + store = CredentialCooldowns(self.path) + store.note_credential(IDENTITY, PROFILE, time.time() + 300) + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + self.assertFalse(store.note_model(IDENTITY, PROFILE, MODEL, time.time() + 600)) + self.assertIsNotNone(store.last_error) + # The in-memory view still reflects the change even though it was not durable. + self.assertGreater(store.model_until(IDENTITY, PROFILE, MODEL), time.time()) + self.assertEqual(json.loads(self.path.read_text())["accounts"][IDENTITY]["models"], {}) + + def test_a_failed_clear_is_not_reported_as_durable(self): + store = CredentialCooldowns(self.path) + store.note_credential(IDENTITY, PROFILE, time.time() + 300) + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + outcome = store.clear_credential(IDENTITY, PROFILE) + self.assertTrue(outcome["changed"]) # The in-memory row was dropped. + self.assertFalse(outcome["durable"]) # But the disk row survived. + self.assertIsNotNone(store.last_error) + self.assertEqual(store.credential_until(IDENTITY, PROFILE), 0.0) # Cleared in memory. + self.assertGreater(json.loads(self.path.read_text())["accounts"][IDENTITY]["fail_until"], 0) + + def test_a_clear_retries_a_previously_failed_write(self): + """A second clear must retry the write, not report "nothing to do" while disk is stale.""" + store = CredentialCooldowns(self.path) + store.note_credential(IDENTITY, PROFILE, time.time() + 300) + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + self.assertFalse(store.clear_credential(IDENTITY, PROFILE)["durable"]) + self.assertGreater(json.loads(self.path.read_text())["accounts"][IDENTITY]["fail_until"], 0) + # The retry runs with the fault removed. + outcome = store.clear_credential(IDENTITY, PROFILE) + self.assertFalse(outcome["changed"]) # Nothing left to change in memory. + self.assertTrue(outcome["durable"]) # But the stale disk row was cleared. + self.assertNotIn(IDENTITY, json.loads(self.path.read_text())["accounts"]) + + def test_clear_outcomes_distinguish_noop_from_failure(self): + store = CredentialCooldowns(self.path) + # Nothing recorded at all: a genuine no-op, which is durable by definition. + self.assertEqual(store.clear_credential(IDENTITY, PROFILE), {"changed": False, "durable": True}) + self.assertEqual(store.clear_model(IDENTITY, PROFILE, MODEL), {"changed": False, "durable": True}) + self.assertEqual(store.forget(IDENTITY), {"changed": False, "durable": True}) + + def test_missing_path_stays_in_memory_only(self): + store = CredentialCooldowns() + self.assertTrue(store.note_credential(IDENTITY, PROFILE, time.time() + 300)) + self.assertGreater(store.credential_until(IDENTITY, PROFILE), time.time()) + self.assertEqual(list(self.root.iterdir()), []) + + def test_a_symlinked_file_is_not_adopted_and_the_target_is_untouched(self): + target = self.root / "outside.json" + target.write_text(json.dumps({"version": VERSION, "accounts": {IDENTITY: self.row()}}), encoding="utf-8") + link = self.root / "link.json" + try: + link.symlink_to(target) + except (OSError, NotImplementedError): + self.skipTest("symlinks are unavailable on this platform") + self.assertEqual(CredentialCooldowns(link).detail(), []) + self.assertTrue(target.exists()) + + def test_a_fifo_is_not_opened(self): + if sys.platform == "win32": + self.skipTest("FIFOs are a POSIX feature") + fifo = self.root / "pipe.json" + os.mkfifo(fifo) + self.assertEqual(CredentialCooldowns(fifo).detail(), []) # Must not block. + + def test_snapshot_is_owner_only_where_the_platform_supports_it(self): + store = CredentialCooldowns(self.path) + store.note_credential(IDENTITY, PROFILE, time.time() + 300) + if sys.platform != "win32": + self.assertEqual(self.path.stat().st_mode & 0o777, 0o600) + + +class PoolCooldownPersistenceTests(unittest.TestCase): + """Verify the pool adopts and records cooldowns through its real public methods.""" + + def setUp(self): + self.root = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": str(self.root)})) + self.enterContext(patch.dict(converter.CONFIG, {"log_path": None, "cred_pool": None, "cred": None, + "ledger": None, "model_catalogs": {}, + "account_catalogs": None, "model_cache": None})) + self.enterContext(patch.object(converter, "_log")) + self.path = self.root / "credential-cooldowns.json" + + def credential(self, name="account.info", uid="synthetic-uid", domain="www.codebuddy.cn", + access_token="synthetic-token"): + now = time.time() + path = self.root / name + path.write_text(json.dumps({"account": {"uid": uid}, "auth": { + "accessToken": access_token, "refreshToken": "synthetic-refresh", "domain": domain, + "expiresAt": (now + 86400) * 1000, "lastRefreshTime": now * 1000}}), encoding="utf-8") + return path + + def pool(self, *paths): + return converter.CredentialPool(list(paths), blocks_path=self.root / "blocks.json", + cooldowns_path=self.path) + + def test_auth_circuit_breaker_survives_a_restart(self): + first = self.pool(self.credential()) + first.cooldown(first._entries[0]["cm"], reason="backend HTTP 401") + self.assertFalse(first._healthy(first._entries[0])) + revived = self.pool(self.credential()) + self.assertGreater(revived._entries[0]["fail_until"], time.time()) + self.assertEqual(revived._entries[0]["last_error"], "backend HTTP 401") + self.assertFalse(revived._healthy(revived._entries[0])) + + def test_model_cooldown_survives_a_restart_and_stays_per_model(self): + first = self.pool(self.credential()) + first.note_status(first._entries[0]["cm"], 429, model=MODEL, raw=b"") + revived = self.pool(self.credential()) + self.assertFalse(revived._model_healthy(revived._entries[0], MODEL)) + self.assertTrue(revived._model_healthy(revived._entries[0], "another-model")) + + def test_repeated_restarts_do_not_extend_the_breaker(self): + first = self.pool(self.credential()) + first.cooldown(first._entries[0]["cm"], reason="backend HTTP 401") + original = first._entries[0]["fail_until"] + for _ in range(4): + self.pool(self.credential()) + self.assertLessEqual(self.pool(self.credential())._entries[0]["fail_until"], original + 0.01) + + def test_an_explicit_reset_lifts_both_kinds_of_cooldown(self): + first = self.pool(self.credential()) + entry = first._entries[0] + first.cooldown(entry["cm"], reason="backend HTTP 401") + first.note_status(entry["cm"], 429, model=MODEL, raw=b"") + self.assertTrue(first.clear_cooldowns(entry["cm"], MODEL)["durable"]) + self.assertTrue(first.clear_cooldowns(entry["cm"])["durable"]) + revived = self.pool(self.credential())._entries[0] + self.assertTrue(self.pool(self.credential())._healthy(revived)) + self.assertTrue(self.pool(self.credential())._model_healthy(revived, MODEL)) + self.assertIsNone(revived.get("last_error")) + + def test_a_same_path_replacement_does_not_inherit_the_previous_accounts_cooldowns(self): + """A credential file rewritten for another account must not keep the old cooldowns.""" + path = self.credential(name="shared.info", uid="first-uid") + first = self.pool(path) + first.cooldown(first._entries[0]["cm"], reason="backend HTTP 401") + first.note_status(first._entries[0]["cm"], 429, model=MODEL, raw=b"") + self.assertGreater(self.pool(path)._entries[0]["fail_until"], time.time()) + + # Rewrite the same path for a different account, which already has its own cooldowns. + self.credential(name="shared.info", uid="second-uid") + second = self.pool(path) + replacement = second._entries[0] + self.assertNotEqual(replacement["account_key"], first._entries[0]["account_key"]) + self.assertEqual(replacement["fail_until"], 0.0) + self.assertEqual(second._model_fail, {}) + self.assertTrue(second._model_healthy(replacement, MODEL)) + + def test_in_process_replacement_keeps_the_incoming_accounts_own_breaker(self): + """The same replacement, but through reload() on one live pool. + + This is the path the gateway actually takes when a credential file changes on disk, + and it is where an explicit reset must not wipe the *incoming* account's breaker. + """ + shared = self.credential(name="shared.info", uid="first-uid") + incoming = self.credential(name="incoming.info", uid="second-uid") + # The incoming account's breaker is already on disk, as it would be after a restart. + seed = self.pool(incoming) + seed.cooldown(seed._entries[0]["cm"], reason="backend HTTP 403") + identity = seed._entries[0]["account_key"] + + pool = self.pool(shared, incoming) + self.assertNotEqual(pool._entries[0]["account_key"], identity) + # Point the shared path at the incoming account and reload in place. + self.credential(name="shared.info", uid="second-uid") + pool.reload([shared], reset=True) + entry = next(e for e in pool._entries if Path(e["id"]).name == "shared.info") + self.assertEqual(entry["account_key"], identity) + self.assertGreater(entry["fail_until"], time.time()) # B's breaker survived. + self.assertEqual(entry["last_error"], "backend HTTP 403") # Not A's reason. + self.assertFalse(pool._healthy(entry)) + + def test_an_account_without_a_uid_never_keys_durable_state(self): + """An empty UID still hashes, so every such account would share one identity.""" + path = self.root / "anonymous.info" + now = time.time() + path.write_text(json.dumps({"account": {}, "auth": { + "accessToken": "synthetic-token", "refreshToken": "synthetic-refresh", + "domain": "www.codebuddy.cn", "expiresAt": (now + 86400) * 1000, + "lastRefreshTime": now * 1000}}), encoding="utf-8") + pool = self.pool(path) + entry = pool._entries[0] + self.assertFalse(pool._durable_identity(entry)) + pool.cooldown(entry["cm"], reason="backend HTTP 401") + self.assertFalse(pool._healthy(entry)) # Still enforced in memory. + self.assertEqual(pool.cooldown_detail(), []) # But never written to disk. + + def test_clear_cooldowns_reports_memory_and_durable_separately(self): + pool = self.pool(self.credential()) + entry = pool._entries[0] + pool.cooldown(entry["cm"], reason="backend HTTP 401") + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + outcome = pool.clear_cooldowns(entry["cm"]) + self.assertTrue(outcome["changed_in_memory"]) # The runtime view is reset. + self.assertFalse(outcome["durable"]) # But it is not durable. + self.assertTrue(pool._healthy(entry)) + + def test_a_failed_clear_is_never_reported_as_durable(self): + """Memory already lacking a row says nothing about whether disk still has it.""" + pool = self.pool(self.credential()) + entry = pool._entries[0] + pool.note_status(entry["cm"], 429, model=MODEL, raw=b"") + pool._model_fail.pop((entry["id"], MODEL), None) # Memory is already clear... + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + outcome = pool.clear_cooldowns(entry["cm"], MODEL) + self.assertFalse(outcome["changed_in_memory"]) + self.assertFalse(outcome["durable"]) # ...but the disk row survived. + self.assertFalse(self.pool(entry["cm"].path)._model_healthy( + self.pool(entry["cm"].path)._entries[0], MODEL)) + + def test_a_write_failure_is_reported_operationally_once_per_interval(self): + pool = self.pool(self.credential()) + entry = pool._entries[0] + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + for _ in range(3): + pool.cooldown(entry["cm"], reason="backend HTTP 401") + warnings = [call for call in converter._log.call_args_list + if "持久化失败" in str(call.args[0] if call.args else "")] + self.assertEqual(len(warnings), 1) # Rate limited, not per failure. + + def test_a_replacement_inherits_only_its_own_persisted_cooldowns(self): + old = self.credential(name="shared.info", uid="first-uid") + first = self.pool(old) + first.cooldown(first._entries[0]["cm"], reason="backend HTTP 401") + # Give the incoming account its own persisted breaker before the swap. + incoming = self.credential(name="incoming.info", uid="second-uid") + probe = self.pool(incoming) + probe.cooldown(probe._entries[0]["cm"], reason="backend HTTP 403") + identity = probe._entries[0]["account_key"] + self.credential(name="shared.info", uid="second-uid") + revived = self.pool(old) + self.assertEqual(revived._entries[0]["account_key"], identity) + self.assertGreater(revived._entries[0]["fail_until"], time.time()) + self.assertEqual(revived._entries[0]["last_error"], "backend HTTP 403") + + def test_an_explicit_reload_lifts_the_breaker_and_keeps_model_cooldowns(self): + path = self.credential() + first = self.pool(path) + entry = first._entries[0] + first.cooldown(entry["cm"], reason="backend HTTP 401") + first.note_status(entry["cm"], 429, model=MODEL, raw=b"") + first.reload([path], reset=True) # The upstream explicit-reload path. + self.assertTrue(first._healthy(entry)) + self.assertFalse(first._model_healthy(entry, MODEL)) + revived = self.pool(path) # A restart must not resurrect the breaker. + self.assertTrue(revived._healthy(revived._entries[0])) + self.assertFalse(revived._model_healthy(revived._entries[0], MODEL)) + + def test_a_same_account_token_refresh_keeps_cooldowns(self): + """A refresh writes a new token for the same account; the 429 cooldown must persist.""" + path = self.credential() + first = self.pool(path) + entry = first._entries[0] + first.note_status(entry["cm"], 429, model=MODEL, raw=b"") + generation = entry["generation"] + # A real refresh: the token file is rewritten and the manager is invalidated, which is + # what CredentialManager itself does after writing new credentials. + self.credential(access_token="rotated-access-token") + entry["cm"].invalidate() + first.reload([path], reset=False) + self.assertNotEqual(entry["generation"], generation) # The reload really re-read it. + self.assertEqual(entry["account_key"], first._entries[0]["account_key"]) + self.assertFalse(first._model_healthy(entry, MODEL)) + # And it survives a restart too. + self.assertFalse(self.pool(path)._model_healthy(self.pool(path)._entries[0], MODEL)) + + def test_deleting_and_re_adding_a_credential_starts_clean(self): + path = self.credential() + first = self.pool(path) + first.cooldown(first._entries[0]["cm"], reason="backend HTTP 401") + self.assertTrue(first.remove_file("account.info")) + self.assertEqual(first.cooldown_detail(), []) + revived = self.pool(self.credential()) + self.assertEqual(revived._entries[0]["fail_until"], 0.0) + + def test_international_auto_is_isolated_per_routed_model(self): + """`auto` maps to `default-model` internationally, so the two must not alias.""" + path = self.credential(domain="www.codebuddy.ai") + pool = self.pool(path) + entry = pool._entries[0] + self.assertEqual(entry["profile"], "intl-cli") + pool.note_status(entry["cm"], 429, model="auto", raw=b"") + revived = self.pool(path) + self.assertFalse(revived._model_healthy(revived._entries[0], "auto")) + self.assertIn(("default-model"), [model for _, model in revived._model_fail]) + + def test_a_pool_without_a_path_keeps_cooldowns_in_memory(self): + pool = converter.CredentialPool([self.credential()], blocks_path=self.root / "blocks.json") + pool.cooldown(pool._entries[0]["cm"], reason="backend HTTP 401") + self.assertFalse(pool._healthy(pool._entries[0])) + self.assertEqual(list(self.root.glob("credential-cooldowns.json")), []) + + def test_a_write_failure_is_visible_and_still_degrades_gracefully(self): + pool = self.pool(self.credential()) + entry = pool._entries[0] + with patch("app.credential_cooldowns.os.replace", side_effect=OSError("read-only")): + pool.cooldown(entry["cm"], reason="backend HTTP 401") + self.assertFalse(pool._healthy(entry)) # Runtime behaviour is unchanged. + self.assertTrue(pool.cooldown_storage()["degraded"]) + self.assertIsNotNone(pool.cooldown_storage()["last_error"]) + + def test_cooldowns_are_visible_to_the_admin_snapshot(self): + pool = self.pool(self.credential()) + entry = pool._entries[0] + pool.cooldown(entry["cm"], reason="backend HTTP 401") + pool.note_status(entry["cm"], 429, model=MODEL, raw=b"") + row = self.pool(self.credential()).snapshot()[0] # A restarted pool must still report both. + self.assertFalse(row["healthy"]) + self.assertIn(MODEL, row["model_cooldowns"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime_endpoints.py b/tests/test_runtime_endpoints.py index a7133d0..3b8dab5 100644 --- a/tests/test_runtime_endpoints.py +++ b/tests/test_runtime_endpoints.py @@ -10,6 +10,7 @@ import json import os import tempfile +import time import unittest from concurrent.futures import ThreadPoolExecutor from unittest.mock import Mock, patch @@ -22,6 +23,8 @@ import converter from app import upstream_io +from app import runtime_management +from app.usage_snapshots import UsageSnapshots ROUTES = ("/v1/chat/completions", "/v1/responses", "/v1/messages") @@ -778,6 +781,120 @@ async def test_saturated_generation_gate_does_not_block_token_counting(self): +class StartupUsageHydrationTests(unittest.TestCase): + """Exercise the real main() startup path, not a look-alike helper call.""" + + def setUp(self): + self.root = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": str(self.root)}, clear=False)) + self.path = self.root / "usage-snapshots.json" + + def credential(self, uid="synthetic-uid"): + now = time.time() + path = self.root / "account.info" + path.write_text(json.dumps({"account": {"uid": uid}, "auth": { + "accessToken": "synthetic-token", "refreshToken": "synthetic-refresh", + "domain": "www.codebuddy.cn", "expiresAt": (now + 86400) * 1000, + "lastRefreshTime": now * 1000}}), encoding="utf-8") + return path + + def test_main_publishes_cached_usage_before_serving(self): + """The aggregate must be populated by the time uvicorn.run is entered.""" + path = self.credential() + # Preseed the cache for the identity this credential will resolve to. + probe = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = probe._entries[0] + snapshots = UsageSnapshots(self.path) + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {"2026-09-19": {"m": 6.0}}, "total_credits": 6.0, + "requests": 3, "partial": False}) + + observed = {} + with contextlib.ExitStack() as stack: + stack.enter_context(patch.dict(converter.CONFIG)) + stack.enter_context(patch("sys.argv", ["converter.py", "--skip-check", + "--auth-file", str(path)])) + stack.enter_context(contextlib.redirect_stderr(io.StringIO())) + stack.enter_context(patch.object(converter.threading, "Thread")) # No real threads. + stack.enter_context(patch.object(converter, "credits_mod", None)) # No network. + stack.enter_context(patch.object(converter, "seed_credentials")) + stack.enter_context(patch.object(converter, "preflight", return_value=True)) + # Keep this test on startup ordering: skip the SQLite-backed stores so no file + # handle is left open on Windows. + stack.enter_context(patch.object(runtime_management, "initialize")) + stack.enter_context(patch.object(runtime_management, "install")) + + def capture(*args, **kwargs): + # Inspect the aggregate at the moment the server would start serving. + observed["usage"] = converter.CONFIG.get("usage_daily") + observed["accounts"] = converter.CONFIG.get("usage_daily_accounts") + stack.enter_context(patch.object(converter.uvicorn, "run", side_effect=capture)) + converter.main() + + self.assertIsNotNone(observed["usage"], "main() never published usage") + self.assertEqual(observed["usage"]["total_credits"], 6.0) + self.assertEqual(observed["usage"]["requests"], 3) + self.assertEqual(observed["usage"]["stale_accounts"], ["account.info"]) + self.assertTrue(observed["accounts"][entry["id"]]["stale"]) + + def test_main_without_a_cache_leaves_usage_empty(self): + """An empty deployment must not be disturbed by the hydration step.""" + path = self.credential() + observed = {} + with contextlib.ExitStack() as stack: + stack.enter_context(patch.dict(converter.CONFIG)) + stack.enter_context(patch("sys.argv", ["converter.py", "--skip-check", + "--auth-file", str(path)])) + stack.enter_context(contextlib.redirect_stderr(io.StringIO())) + stack.enter_context(patch.object(converter.threading, "Thread")) + stack.enter_context(patch.object(converter, "credits_mod", None)) + stack.enter_context(patch.object(converter, "seed_credentials")) + stack.enter_context(patch.object(converter, "preflight", return_value=True)) + # Keep this test on startup ordering: skip the SQLite-backed stores so no file + # handle is left open on Windows. + stack.enter_context(patch.object(runtime_management, "initialize")) + stack.enter_context(patch.object(runtime_management, "install")) + stack.enter_context(patch.object(converter.uvicorn, "run", + side_effect=lambda *a, **k: observed.setdefault("usage", + converter.CONFIG.get("usage_daily")))) + converter.main() + self.assertIsNone(observed["usage"]) + self.assertFalse(self.path.exists()) + + def test_main_starts_no_maintenance_thread_before_hydration(self): + """Hydration must complete before the maintenance threads are started.""" + path = self.credential() + probe = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = probe._entries[0] + UsageSnapshots(self.path).store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 1.0, "requests": 1, + "partial": False}) + events = [] + with contextlib.ExitStack() as stack: + stack.enter_context(patch.dict(converter.CONFIG)) + stack.enter_context(patch("sys.argv", ["converter.py", "--skip-check", + "--auth-file", str(path)])) + stack.enter_context(contextlib.redirect_stderr(io.StringIO())) + stack.enter_context(patch.object(converter, "credits_mod", None)) + stack.enter_context(patch.object(converter, "seed_credentials")) + stack.enter_context(patch.object(converter, "preflight", return_value=True)) + # Keep this test on startup ordering: skip the SQLite-backed stores so no file + # handle is left open on Windows. + stack.enter_context(patch.object(runtime_management, "initialize")) + stack.enter_context(patch.object(runtime_management, "install")) + real_thread = converter.threading.Thread + + def note_thread(*args, **kwargs): + events.append(("thread", converter.CONFIG.get("usage_daily") is not None)) + return Mock() + stack.enter_context(patch.object(converter.threading, "Thread", side_effect=note_thread)) + stack.enter_context(patch.object(converter.uvicorn, "run")) + converter.main() + # Every thread must have been started after usage was published. + self.assertTrue(events) + self.assertTrue(all(published for _, published in events), events) + + class ConfigurationTests(unittest.TestCase): def configure(self, env=None, flags=(), invalid=False, stored=None, expected_host=None): with contextlib.ExitStack() as stack: diff --git a/tests/test_usage_snapshots.py b/tests/test_usage_snapshots.py new file mode 100644 index 0000000..01ccff2 --- /dev/null +++ b/tests/test_usage_snapshots.py @@ -0,0 +1,591 @@ +#!/usr/bin/env python3 +"""Test the usage-snapshot cache: restart persistence, staleness and identity safety.""" + +import json +import os +import sys +import tempfile +import time +import unittest +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +import converter +from app.usage_snapshots import (MAX_ACCOUNTS, MAX_BYTES, MAX_DAYS, MAX_MODELS, VERSION, + UsageSnapshots) + +IDENTITY = "a" * 64 +OTHER_IDENTITY = "b" * 64 +PATH = "/auth/account.info" +USAGE = {"by_day": {"2026-09-19": {"claude-sonnet-4": 12.5}}, "total_credits": 12.5, + "requests": 7, "partial": False} + + +class UsageSnapshotStoreTests(unittest.TestCase): + """Exercise the cache directly, with no credential files and no network access.""" + + def setUp(self): + self.root = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.path = self.root / "usage-snapshots.json" + + def row(self, **overrides): + row = {"identity": IDENTITY, "site": "domestic", "by_day": {"2026-09-19": {"m": 1.0}}, + "total_credits": 1.0, "requests": 1, "partial": False, "fetched_at": time.time()} + row.update(overrides) + return row + + def write(self, accounts): + self.path.write_text(json.dumps({"version": VERSION, "accounts": accounts}), encoding="utf-8") + + def test_a_snapshot_survives_a_restart(self): + first = UsageSnapshots(self.path) + self.assertTrue(first.store(PATH, IDENTITY, "domestic", USAGE)) + restored = UsageSnapshots(self.path).accounts()[PATH] + self.assertEqual(restored["total_credits"], 12.5) + self.assertEqual(restored["by_day"]["2026-09-19"]["claude-sonnet-4"], 12.5) + self.assertEqual(restored["site"], "domestic") + + def test_a_cached_row_keeps_its_identity_and_staleness_flags(self): + store = UsageSnapshots(self.path) + store.store(PATH, IDENTITY, "international", USAGE, partial=True) + row = UsageSnapshots(self.path).accounts()[PATH] + self.assertEqual(row["identity"], IDENTITY) + self.assertTrue(row["partial"]) + + def test_identity_mismatch_is_detectable(self): + store = UsageSnapshots(self.path) + store.store(PATH, IDENTITY, "domestic", USAGE) + revived = UsageSnapshots(self.path) + self.assertTrue(revived.identity_matches(PATH, IDENTITY)) + self.assertFalse(revived.identity_matches(PATH, OTHER_IDENTITY)) + self.assertFalse(revived.identity_matches("/auth/unknown.info", IDENTITY)) + + def test_untrusted_snapshots_are_rejected_wholesale(self): + payloads = { + "not json": b"nope", + "empty object": b"{}", + "unsupported version": json.dumps({"version": 99, "accounts": {}}).encode(), + "boolean version": json.dumps({"version": True, "accounts": {}}).encode(), + "extra top-level field": json.dumps({"version": VERSION, "accounts": {}, "x": 1}).encode(), + "extra row field": json.dumps({"version": VERSION, + "accounts": {PATH: {**self.row(), "x": 1}}}).encode(), + "non-hex identity": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(identity="zzz")}}).encode(), + "unknown site": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(site="mars")}}).encode(), + "duplicate keys": b'{"version":1,"version":1,"accounts":{}}', + "infinity literal": ('{"version":1,"accounts":{"' + PATH + '":{"identity":"' + IDENTITY + + '","site":"domestic","by_day":{},"total_credits":Infinity,' + '"requests":0,"partial":false,"fetched_at":1}}}').encode(), + "huge integer": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(total_credits=10 ** 400)}}).encode(), + "negative credits": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(total_credits=-5.0)}}).encode(), + "bool partial": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(partial=1)}}).encode(), + "bad date": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(by_day={"19-09-2026": {"m": 1.0}})}}).encode(), + "impossible date": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(by_day={"2026-02-31": {"m": 1.0}})}}).encode(), + "fractional requests": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(requests=1.5)}}).encode(), + "future timestamp": json.dumps({"version": VERSION, + "accounts": {PATH: self.row(fetched_at=time.time() + 10 * 86400)}}).encode(), + "oversized file": b"x" * (MAX_BYTES + 1), + } + for label, content in payloads.items(): + with self.subTest(payload=label): + self.path.write_bytes(content) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + + def test_oversized_tables_are_rejected_on_load(self): + accounts = {f"/auth/{index}.info": self.row() for index in range(MAX_ACCOUNTS + 1)} + self.write(accounts) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + self.write({PATH: self.row(by_day={f"2026-09-{day:02d}": {} for day in range(1, MAX_DAYS + 2)})}) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + self.write({PATH: self.row(by_day={"2026-09-19": {f"m{index}": 1.0 + for index in range(MAX_MODELS + 1)}})}) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + + def test_the_writer_never_produces_a_file_the_reader_rejects(self): + """A rejected cache file loses every snapshot it held, not just the offending row.""" + store = UsageSnapshots(self.path) + # October has 31 days, matching MAX_DAYS, so the fixtures are real calendar dates. + days = [f"2026-10-{day:02d}" for day in range(1, MAX_DAYS + 1)] + with store._lock: + for index in range(MAX_ACCOUNTS): + store._data[f"/auth/{index}.info"] = { + "identity": f"{index:064x}", "site": "domestic", + "by_day": {day: {f"model-name-{model:03d}": 1.5 for model in range(MAX_MODELS)} + for day in days}, + "total_credits": 1.5, "requests": 1, "partial": False, + "fetched_at": time.time() - index} + content, shed = store._serialize_locked() + self.assertLessEqual(len(content), MAX_BYTES) + self.assertTrue(shed) # Capacity really did force shedding. + self.path.write_bytes(content) + self.assertGreater(len(UsageSnapshots(self.path).accounts()), 0) + + def test_a_write_failure_is_reported_and_keeps_the_cache_in_memory(self): + store = UsageSnapshots(self.path) + store.store(PATH, IDENTITY, "domestic", USAGE) + with patch("app.usage_snapshots.os.replace", side_effect=OSError("read-only")): + self.assertFalse(store.store("/auth/other.info", IDENTITY, "domestic", USAGE)) + self.assertIsNotNone(store.last_error) + self.assertIn("/auth/other.info", store.accounts()) # Still usable in memory. + self.assertNotIn("/auth/other.info", json.loads(self.path.read_text())["accounts"]) + + def test_prune_drops_stale_rows_only(self): + store = UsageSnapshots(self.path) + store.store("/auth/old.info", IDENTITY, "domestic", USAGE) + store.store("/auth/new.info", IDENTITY, "domestic", USAGE) + with store._lock: + store._data["/auth/old.info"]["fetched_at"] = time.time() - 30 * 24 * 3600 + self.assertTrue(store.prune()) + self.assertEqual(set(UsageSnapshots(self.path).accounts()), {"/auth/new.info"}) + + def test_forget_removes_one_snapshot(self): + store = UsageSnapshots(self.path) + store.store(PATH, IDENTITY, "domestic", USAGE) + self.assertTrue(store.forget(PATH)) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + + def test_invalid_input_is_refused_rather_than_cached(self): + store = UsageSnapshots(self.path) + self.assertFalse(store.store(PATH, "not-a-hash", "domestic", USAGE)) + self.assertFalse(store.store(PATH, IDENTITY, "mars", USAGE)) + self.assertFalse(store.store(PATH, IDENTITY, "domestic", {"by_day": {"nope": {"m": 1.0}}})) + self.assertEqual(store.accounts(), {}) + + def test_malformed_usage_is_rejected_rather_than_normalized_to_zero(self): + """A missing or bad field must not become a plausible zero.""" + store = UsageSnapshots(self.path) + bad = { + "missing by_day": {"total_credits": 1.0, "requests": 1}, + "missing total": {"by_day": {}, "requests": 1}, + "missing requests": {"by_day": {}, "total_credits": 1.0}, + "none by_day": {"by_day": None, "total_credits": 1.0, "requests": 1}, + "none total": {"by_day": {}, "total_credits": None, "requests": 1}, + "empty-string total": {"by_day": {}, "total_credits": "", "requests": 1}, + "false total": {"by_day": {}, "total_credits": False, "requests": 1}, + "string total": {"by_day": {}, "total_credits": "1.0", "requests": 1}, + "negative total": {"by_day": {}, "total_credits": -1.0, "requests": 1}, + "huge total": {"by_day": {}, "total_credits": 10 ** 400, "requests": 1}, + "nan total": {"by_day": {}, "total_credits": float("nan"), "requests": 1}, + "fractional requests": {"by_day": {}, "total_credits": 1.0, "requests": 1.5}, + "false requests": {"by_day": {}, "total_credits": 1.0, "requests": False}, + "negative credit day": {"by_day": {"2026-09-19": {"m": -1.0}}, "total_credits": 1.0, + "requests": 1}, + "string day credit": {"by_day": {"2026-09-19": {"m": "1"}}, "total_credits": 1.0, + "requests": 1}, + "impossible date": {"by_day": {"2026-02-31": {"m": 1.0}}, "total_credits": 1.0, + "requests": 1}, + } + for label, usage in bad.items(): + with self.subTest(usage=label): + self.assertFalse(store.store(PATH, IDENTITY, "domestic", usage)) + self.assertEqual(store.accounts(), {}) + + def test_non_mapping_usage_and_partial_flag_are_rejected(self): + store = UsageSnapshots(self.path) + self.assertFalse(store.store(PATH, IDENTITY, "domestic", None)) + self.assertFalse(store.store(PATH, IDENTITY, "domestic", ["by_day"])) + for partial in (1, 0, "true", None): + with self.subTest(partial=partial): + self.assertFalse(store.store(PATH, IDENTITY, "domestic", USAGE, partial=partial)) + self.assertEqual(store.accounts(), {}) + + def test_the_stored_row_is_a_copy_the_caller_cannot_mutate(self): + store = UsageSnapshots(self.path) + usage = {"by_day": {"2026-09-19": {"m": 1.0}}, "total_credits": 1.0, "requests": 1, + "partial": False} + self.assertTrue(store.store(PATH, IDENTITY, "domestic", usage)) + usage["by_day"]["2026-09-19"]["m"] = 999.0 # Mutating the caller's object... + usage["total_credits"] = 999.0 + self.assertEqual(store.accounts()[PATH]["total_credits"], 1.0) # ...must not leak in. + self.assertEqual(store.accounts()[PATH]["by_day"]["2026-09-19"]["m"], 1.0) + + def test_an_overlong_path_is_rejected_on_write_and_load(self): + store = UsageSnapshots(self.path) + self.assertFalse(store.store("/" + "x" * 5000, IDENTITY, "domestic", USAGE)) + self.write({"/" + "x" * 5000: self.row()}) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + + def test_missing_path_stays_in_memory_only(self): + store = UsageSnapshots() + self.assertTrue(store.store(PATH, IDENTITY, "domestic", USAGE)) + self.assertIn(PATH, store.accounts()) + self.assertEqual(list(self.root.iterdir()), []) + + def test_a_symlinked_file_is_not_adopted(self): + target = self.root / "outside.json" + target.write_text(json.dumps({"version": VERSION, "accounts": {PATH: self.row()}}), encoding="utf-8") + link = self.root / "link.json" + try: + link.symlink_to(target) + except (OSError, NotImplementedError): + self.skipTest("symlinks are unavailable on this platform") + self.assertEqual(UsageSnapshots(link).accounts(), {}) + self.assertTrue(target.exists()) + + def test_snapshot_is_owner_only_where_the_platform_supports_it(self): + UsageSnapshots(self.path).store(PATH, IDENTITY, "domestic", USAGE) + if sys.platform != "win32": + self.assertEqual(self.path.stat().st_mode & 0o777, 0o600) + + +class UsageSnapshotIntegrationTests(unittest.TestCase): + """Verify the cache feeds the published aggregate without becoming authoritative.""" + + def setUp(self): + self.root = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.enterContext(patch.dict(os.environ, {"CODEBUDDY_AUTH_DIR": str(self.root)})) + self.enterContext(patch.dict(converter.CONFIG, {"log_path": None, "cred_pool": None, "cred": None, + "ledger": None, "model_catalogs": {}, + "account_catalogs": None, "model_cache": None, + "usage_daily": None, "usage_daily_accounts": None, + "usage_snapshots": None})) + self.enterContext(patch.object(converter, "_log")) + self.path = self.root / "usage-snapshots.json" + + def credential(self, name="account.info", uid="synthetic-uid"): + now = time.time() + path = self.root / name + path.write_text(json.dumps({"account": {"uid": uid}, "auth": { + "accessToken": "synthetic-token", "refreshToken": "synthetic-refresh", + "domain": "www.codebuddy.cn", "expiresAt": (now + 86400) * 1000, + "lastRefreshTime": now * 1000}}), encoding="utf-8") + return path + + def test_the_published_aggregate_is_repopulated_from_the_cache(self): + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {"2026-09-19": {"claude-sonnet-4": 4.0}}, "total_credits": 4.0, + "requests": 3, "partial": False}) + # A restart starts with an empty in-memory aggregate. + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(pool) + published = converter.CONFIG["usage_daily"] + self.assertEqual(published["total_credits"], 4.0) + self.assertEqual(published["requests"], 3) + self.assertEqual(published["by_day"]["2026-09-19"]["claude-sonnet-4"], 4.0) + + def test_a_reused_path_does_not_show_the_previous_accounts_usage(self): + path = self.credential(name="shared.info", uid="first-uid") + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + snapshots.store(pool._entries[0]["id"], pool._entries[0]["account_key"], "domestic", + {"by_day": {"2026-09-19": {"m": 9.0}}, "total_credits": 9.0, + "requests": 9, "partial": False}) + # Rewrite the same path for another account. + self.credential(name="shared.info", uid="second-uid") + replacement = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(replacement) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 0.0) + + def test_a_live_snapshot_is_not_overwritten_by_the_cache(self): + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 1.0, "requests": 1, "partial": False}) + converter.CONFIG["usage_daily_accounts"] = {entry["id"]: { + "identity": entry["account_key"], "site": "domestic", "by_day": {}, + "total_credits": 99.0, "requests": 99, + "partial": False, "fetched_at": time.time()}} + converter._publish_usage_daily(pool) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 99.0) + + def test_deleting_a_credential_forgets_its_cached_usage(self): + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 1.0, "requests": 1, "partial": False}) + converter.CONFIG["usage_daily_accounts"] = {entry["id"]: { + "identity": entry["account_key"], "site": "domestic", "by_day": {}, + "total_credits": 1.0, "requests": 1, "partial": False, "fetched_at": time.time()}} + self.assertTrue(pool.remove_file("account.info")) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + # The live aggregate row must go too, or a re-add would resurrect the old figure. + self.assertNotIn(entry["id"], converter.CONFIG["usage_daily_accounts"]) + + def test_re_adding_a_deleted_credential_starts_with_no_usage(self): + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + converter.CONFIG["usage_snapshots"] = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"].store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 5.0, "requests": 5, + "partial": False}) + converter.CONFIG["usage_daily_accounts"] = {entry["id"]: { + "identity": entry["account_key"], "site": "domestic", "by_day": {}, + "total_credits": 5.0, "requests": 5, "partial": False, "fetched_at": time.time()}} + pool.remove_file("account.info") + # Re-add the same path and re-publish from a cold aggregate. + self.credential() + revived = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + converter._publish_usage_daily(revived) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 0.0) + + def test_deleting_a_credential_also_drops_its_persisted_cooldowns(self): + """The two cleanups are independent; deletion must run both.""" + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json", + cooldowns_path=self.root / "credential-cooldowns.json") + entry = pool._entries[0] + converter.CONFIG["usage_snapshots"] = UsageSnapshots(self.path) + pool.cooldown(entry["cm"], reason="backend HTTP 401") + converter.CONFIG["usage_snapshots"].store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 1.0, "requests": 1, + "partial": False}) + self.assertTrue(pool.remove_file("account.info")) + self.assertEqual(pool.cooldown_detail(), []) + self.assertEqual(UsageSnapshots(self.path).accounts(), {}) + + def test_a_write_failure_during_a_real_sync_still_publishes_usage(self): + """Drive the real maintenance path, so the patched writer is genuinely exercised.""" + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + converter.CONFIG["usage_daily_accounts"] = None + usage = {"by_day": {"2026-09-19": {"m": 3.0}}, "total_credits": 3.0, "requests": 2, + "partial": False} + with patch.object(converter.credits_mod, "fetch_request_usage", return_value=usage), patch("app.usage_snapshots.os.replace", side_effect=OSError("read-only")) as writer: + stale = converter._sync_usage(pool) + self.assertTrue(writer.called, "the cache writer was never reached") + self.assertEqual(stale, set()) + # The failure is recorded, but the dashboard still shows the freshly fetched figures. + self.assertIsNotNone(snapshots.last_error) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 3.0) + self.assertNotIn("stale_accounts", converter.CONFIG["usage_daily"]) + + def test_a_concurrent_deletion_during_publication_does_not_raise(self): + """Publication must aggregate a consistent snapshot, not index the live map. + + forget_usage() pops from the same dictionary under the pool lock. A reader that copies the + keys and then indexes the live dict therefore raises KeyError when a credential is deleted + mid-pass, which aborts the usage pass. The live rows are seeded directly so the hook below + fires only in the aggregation loop, and the lock is asserted rather than timed. + """ + first = self.credential(name="first.info", uid="first-uid") + second = self.credential(name="second.info", uid="second-uid") + pool = converter.CredentialPool([first, second], blocks_path=self.root / "blocks.json") + rows = {} + for entry in pool._entries: + rows[entry["id"]] = {"identity": entry["account_key"], "site": "domestic", + "by_day": {"2026-09-19": {"m": 3.0}}, "total_credits": 3.0, + "requests": 3, "partial": False, "fetched_at": time.time()} + converter.CONFIG["usage_snapshots"] = UsageSnapshots(self.path) # Empty store. + converter.CONFIG["usage_daily_accounts"] = dict(rows) + victim = pool._entries[1]["id"] + + original = converter._usage_row_expired + observed = [] + + def expiring(snap): + if not observed: + observed.append(True) + # Runs inside the reader's aggregation loop: assert mutual exclusion, then delete + # the row the reader has not reached yet. + observed.append(pool._lock._is_owned()) + converter.CONFIG["usage_daily_accounts"].pop(victim, None) + return original(snap) + + with patch.object(converter, "_usage_row_expired", side_effect=expiring): + converter._publish_usage_daily(pool) # Must not raise KeyError. + self.assertEqual(observed[:2], [True, True], "interleaving not reached, or read unlocked") + self.assertIsNotNone(converter.CONFIG["usage_daily"]) + self.assertNotIn(victim, converter.CONFIG["usage_daily_accounts"]) + # The pass aggregates the snapshot it took at entry, so a row deleted mid-pass is still + # counted this once and simply disappears on the next pass. Consistency is the contract. + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 6.0) + converter._publish_usage_daily(pool) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 3.0) + + def test_adoption_holds_the_pool_lock_while_reading_the_map(self): + """_adopt_cached_usage() reads and prunes the live map, so it needs the same lock.""" + path = self.credential(name="shared.info", uid="first-uid") + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + snapshots = UsageSnapshots(self.path) + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {"2026-09-19": {"m": 7.0}}, "total_credits": 7.0, + "requests": 7, "partial": False}) + converter.CONFIG["usage_snapshots"] = snapshots + # A live row whose recorded owner no longer matches forces the prune path. + converter.CONFIG["usage_daily_accounts"] = {entry["id"]: { + "identity": "0" * 64, "site": "domestic", "by_day": {}, + "total_credits": 1.0, "requests": 1, "fetched_at": time.time()}} + + original = UsageSnapshots.accounts + observed = [] + + def accounts_probe(self): + rows = original(self) + if not observed: + observed.append(pool._lock._is_owned()) + return rows + + with patch.object(UsageSnapshots, "accounts", accounts_probe): + converter._publish_usage_daily(pool) + self.assertEqual(observed, [True], "the live map was read without the pool lock") + # The mismatched row was pruned, and the cache refilled it for the current account. + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 7.0) + + def test_a_malformed_partial_flag_is_rejected_not_coerced(self): + """A non-boolean partial must never be masked into a legitimate-looking flag. + + bool("false") is True, so coercing before the store sees the value would persist a + wrong completeness flag for durable state. The raw value is passed through instead. + """ + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + converter.CONFIG["usage_daily_accounts"] = None + for malformed in ("false", "true", 1, 0, None, [], {}): + with self.subTest(partial=malformed): + snapshots._data.clear() + self.path.unlink(missing_ok=True) + usage = {"by_day": {"2026-09-19": {"m": 2.0}}, "total_credits": 2.0, + "requests": 1, "partial": malformed} + with patch.object(converter.credits_mod, "fetch_request_usage", return_value=usage): + converter._sync_usage(pool) + # Nothing durable was written for this account... + self.assertEqual(snapshots.accounts(), {}, malformed) + # ...while the live dashboard row still reports the fetched figures. + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 2.0) + + def test_a_real_sync_persists_and_survives_a_restart(self): + """The same path with a healthy writer must leave a restorable cache behind.""" + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + converter.CONFIG["usage_snapshots"] = UsageSnapshots(self.path) + converter.CONFIG["usage_daily_accounts"] = None + usage = {"by_day": {"2026-09-19": {"m": 4.0}}, "total_credits": 4.0, "requests": 1, + "partial": False} + with patch.object(converter.credits_mod, "fetch_request_usage", return_value=usage): + converter._sync_usage(pool) + entry = pool._entries[0] + cached = UsageSnapshots(self.path).accounts()[entry["id"]] + self.assertEqual(cached["total_credits"], 4.0) + self.assertEqual(cached["identity"], entry["account_key"]) + # A cold aggregate seeded from disk reports the same figure, marked stale. + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(pool) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 4.0) + self.assertEqual(converter.CONFIG["usage_daily"]["stale_accounts"], ["account.info"]) + + def test_startup_publication_hydrates_from_the_cache(self): + """The startup path must populate usage without waiting for a maintenance pass.""" + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {"2026-09-19": {"m": 2.0}}, "total_credits": 2.0, + "requests": 2, "partial": False}) + # Exactly what startup does, with a cold in-memory aggregate. + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(pool) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 2.0) + # A restored snapshot is stale until a refresh confirms it. + self.assertTrue(converter.CONFIG["usage_daily_accounts"][entry["id"]]["stale"]) + self.assertEqual(converter.CONFIG["usage_daily"]["stale_accounts"], ["account.info"]) + # A restored snapshot is an incomplete view, so the aggregate marks itself partial. + self.assertTrue(converter.CONFIG["usage_daily"]["partial"]) + + def test_a_refresh_clears_only_that_accounts_staleness(self): + first = self.credential(name="a.info", uid="uid-a") + second = self.credential(name="b.info", uid="uid-b") + pool = converter.CredentialPool([first, second], blocks_path=self.root / "blocks.json") + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + for entry in pool.entries(): + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 1.0, "requests": 1, "partial": False}) + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(pool) + self.assertEqual(len(converter.CONFIG["usage_daily"]["stale_accounts"]), 2) + # One account refreshes successfully; the other does not. + refreshed, other = pool.entries() + converter.CONFIG["usage_daily_accounts"] = { + refreshed["id"]: {"identity": refreshed["account_key"], "site": "domestic", "by_day": {}, + "total_credits": 3.0, "requests": 2, "partial": False, + "fetched_at": time.time()}} + converter._publish_usage_daily(pool, {other["id"]}) + self.assertEqual(converter.CONFIG["usage_daily"]["stale_accounts"], ["b.info"]) + # The failed account keeps its last good figure (1.0) alongside the refreshed 3.0. + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 4.0) + + def test_a_stale_restored_row_stops_being_shown_once_it_ages_out(self): + path = self.credential() + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 5.0, "requests": 1, "partial": False}) + with snapshots._lock: + snapshots._data[entry["id"]]["fetched_at"] = time.time() - converter.USAGE_CACHE_MAX_AGE_S - 60 + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(pool) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 0.0) + self.assertEqual(snapshots.accounts(), {}) # The aged row is dropped, not kept. + + def test_an_account_without_a_uid_never_owns_durable_usage(self): + """Two accounts with no UID hash alike, so neither may key a cached snapshot.""" + path = self.root / "anonymous.info" + now = time.time() + path.write_text(json.dumps({"account": {}, "auth": { + "accessToken": "t", "refreshToken": "r", "domain": "www.codebuddy.cn", + "expiresAt": (now + 86400) * 1000, "lastRefreshTime": now * 1000}}), encoding="utf-8") + pool = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + entry = pool._entries[0] + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + # Even a row written directly for that hash is not adopted, because the account cannot + # prove it owns it. + snapshots.store(entry["id"], entry["account_key"], "domestic", + {"by_day": {}, "total_credits": 8.0, "requests": 1, "partial": False}) + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(pool) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 0.0) + + def test_hydrated_rows_are_rechecked_against_current_ownership(self): + """A row hydrated on an earlier pass must not survive a later path reuse.""" + path = self.credential(name="shared.info", uid="first-uid") + first = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + snapshots = UsageSnapshots(self.path) + converter.CONFIG["usage_snapshots"] = snapshots + snapshots.store(first._entries[0]["id"], first._entries[0]["account_key"], "domestic", + {"by_day": {}, "total_credits": 7.0, "requests": 1, "partial": False}) + # Hydrate as a live row, as an earlier publication would have. + converter.CONFIG["usage_daily_accounts"] = None + converter._publish_usage_daily(first) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 7.0) + # The path is then reused by another account. + self.credential(name="shared.info", uid="second-uid") + replacement = converter.CredentialPool([path], blocks_path=self.root / "blocks.json") + converter._publish_usage_daily(replacement) + self.assertEqual(converter.CONFIG["usage_daily"]["total_credits"], 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/web/src/enhancements.test.tsx b/web/src/enhancements.test.tsx index 307015d..f164759 100644 --- a/web/src/enhancements.test.tsx +++ b/web/src/enhancements.test.tsx @@ -169,6 +169,7 @@ describe("scoped credential buttons", () => { ["刷新凭证 one.info", "/credentials/one/refresh"], ["签到 one.info", "/credentials/one/checkin"], ["同步余额 one.info", "/credentials/one/sync"], + ["清除冷却 one.info", "/credentials/one/reset-cooldown"], ["批量签到", "/checkin"], ["同步全部余额", "/sync"], ]) { @@ -179,6 +180,45 @@ describe("scoped credential buttons", () => { expect.objectContaining({ timeout: 300000 }), ); } - expect(post).toHaveBeenCalledTimes(5); + expect(post).toHaveBeenCalledTimes(6); + }); + + it("keeps the cooldown reset available for a disabled account and a failed write", async () => { + // Neither condition may hide the button: a failed durable clear has to stay retryable, and a + // disabled account can still hold a cooldown worth clearing. + vi.mocked(useResource).mockReturnValue({ + data: [ + { id: "one", name: "one.info", enabled: false, profile: "cn-cli", health: "circuit_open" }, + ], + reload: vi.fn(), + loading: false, + error: null, + }); + const post = vi.spyOn(api, "post").mockResolvedValue({ + data: { + results: [ + { + id: "one", + name: "one.info", + ok: false, + changed_in_memory: true, + durable: false, + message: "内存冷却已清除，但写入失败，重启后可能恢复；请稍后重试", + }, + ], + }, + }); + render(); + const button = screen.getByRole("button", { name: "清除冷却 one.info" }); + expect((button as HTMLButtonElement).disabled).toBe(false); + await act(async () => fireEvent.click(button)); + expect(post).toHaveBeenLastCalledWith( + "/credentials/one/reset-cooldown", + undefined, + expect.objectContaining({ timeout: 300000 }), + ); + // The failure is surfaced rather than silently swallowed. + expect(await screen.findByText("未完成")).toBeTruthy(); + expect(screen.getByText(/写入失败/)).toBeTruthy(); }); }); diff --git a/web/src/pages/Credentials.tsx b/web/src/pages/Credentials.tsx index 6511a06..0ed78e9 100644 --- a/web/src/pages/Credentials.tsx +++ b/web/src/pages/Credentials.tsx @@ -116,7 +116,7 @@ export function Credentials() { const [notice, setNotice] = useState(null); const [maintenance, setMaintenance] = useState[]>([]); const maintain = ( - action: "refresh" | "checkin" | "sync" | "travel" | "travel-status", + action: "refresh" | "checkin" | "sync" | "travel" | "travel-status" | "reset-cooldown", credential?: Credential, ) => { if (busy) return; @@ -479,6 +479,16 @@ export function Credentials() { > 同步余额 + {/* Not gated on displayed cooldowns or on enablement: a reset that failed + to persist must stay retryable, and a disabled account can still hold + a stale cooldown worth clearing. */} + {c.trial_supported === true && (