diff --git a/modules/usage_tracker.py b/modules/usage_tracker.py index 14a2770..a1f8c60 100644 --- a/modules/usage_tracker.py +++ b/modules/usage_tracker.py @@ -90,35 +90,58 @@ class UsageTracker: return QUOTA_DEFAULTS["search_daily"] return QUOTA_DEFAULTS["default"][metric] - def _ensure_window(self, metric: QuotaKey) -> Tuple[int, Optional[str], datetime, datetime]: + def _ensure_window( + self, + metric: QuotaKey, + init_if_missing: bool = False, + ) -> Tuple[int, Optional[str], Optional[datetime], Optional[datetime]]: """Returns tuple(count, window_start_iso, window_start_dt, reset_at_dt).""" config = self._get_quota_config(metric) window_hours = config["window_hours"] window_delta = timedelta(hours=window_hours) window_data = self._state["windows"].setdefault(metric, {"count": 0, "window_start": None}) + count = int(window_data.get("count", 0)) window_start_iso = window_data.get("window_start") now = datetime.utcnow() + window_start_dt: Optional[datetime] = None if window_start_iso: try: parsed = datetime.fromisoformat(window_start_iso.replace("Z", "")) window_start_dt = floor_to_hour(parsed) except ValueError: - window_start_dt = floor_to_hour(now) - else: - window_start_dt = floor_to_hour(now) - window_data["window_start"] = frame_iso(window_start_dt) + window_start_dt = None - reset_at_dt = window_start_dt + window_delta - if now >= reset_at_dt and window_data.get("window_start"): - window_data["count"] = 0 - window_start_dt = floor_to_hour(now) + reset_at_dt: Optional[datetime] = None + if window_start_dt: reset_at_dt = window_start_dt + window_delta + if now >= reset_at_dt: + window_data["count"] = 0 + count = 0 + if init_if_missing: + window_start_dt = floor_to_hour(now) + window_data["window_start"] = frame_iso(window_start_dt) + reset_at_dt = window_start_dt + window_delta + else: + window_start_dt = None + window_data["window_start"] = None + reset_at_dt = None + + if window_start_dt is None and init_if_missing: + window_start_dt = floor_to_hour(now) window_data["window_start"] = frame_iso(window_start_dt) + reset_at_dt = window_start_dt + window_delta + elif count == 0 and window_data.get("window_start") and not init_if_missing: + # 没有真实用量时清理遗留的窗口起点,使下一次调用时重新计窗。 + window_data["window_start"] = None + window_start_iso = None + + if window_start_dt and not reset_at_dt: + reset_at_dt = window_start_dt + window_delta return ( - int(window_data.get("count", 0)), - window_data["window_start"], + count, + window_data.get("window_start"), window_start_dt, reset_at_dt, ) @@ -126,7 +149,9 @@ class UsageTracker: def check_and_increment(self, metric: QuotaKey) -> Tuple[bool, Dict[str, str]]: """Check quota and increment if allowed. Returns (allowed, info).""" with self._lock: - count, window_start_iso, window_start_dt, reset_at_dt = self._ensure_window(metric) + count, window_start_iso, window_start_dt, reset_at_dt = self._ensure_window( + metric, init_if_missing=True + ) quota = self._get_quota_config(metric)["limit"] if count >= quota: reset_at_iso = frame_iso(reset_at_dt) @@ -139,8 +164,6 @@ class UsageTracker: new_count = count + 1 self._state["windows"][metric]["count"] = new_count - if not window_start_iso: - self._state["windows"][metric]["window_start"] = frame_iso(window_start_dt) self._state["updated_at"] = self._now_iso() self._save()