JOJO/agent-Specialization
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: delay quota window start until first use

Browse Source
This commit is contained in:
JOJO 2025-11-24 01:38:24 +08:00
parent fd797e3e36
commit 7c7038b5c9
1 changed files with 37 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
modules/usage_tracker.py
View File

@ -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: Optional[datetime] = None
if window_start_dt:
reset_at_dt = window_start_dt + window_delta
if now >= reset_at_dt and window_data.get("window_start"):
if now >= reset_at_dt:
window_data["count"] = 0
count = 0
if init_if_missing:
window_start_dt = floor_to_hour(now)
reset_at_dt = window_start_dt + window_delta
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()