feat: notify on unexpected quota resets

This commit is contained in:
zhoujun0601
2026-08-13 09:51:40 -04:00
parent 63bd0b4fea
commit 984276b359
7 changed files with 291 additions and 35 deletions
+54 -2
View File
@@ -449,14 +449,66 @@ func (a *App) syncAccount(ctx context.Context, id int64) error {
_, _ = a.store.DB.Exec("INSERT INTO daily_usage(account_id,date,total_tokens,fetched_at) VALUES(?,?,?,?) ON CONFLICT(account_id,date) DO UPDATE SET total_tokens=excluded.total_tokens,fetched_at=excluded.fetched_at", id, x.StartDate, x.Tokens, d.FetchedAt)
}
}
for _, x := range d.Limits {
_, _ = a.store.DB.Exec("INSERT INTO limit_snapshots(limit_id,window_type,used_percent,duration_mins,resets_at,fetched_at,account_id) VALUES(?,?,?,?,?,?,?)", x.LimitID, x.WindowType, x.UsedPercent, x.WindowDurationMinutes, x.ResetsAt, d.FetchedAt, id)
resetDetected, e := a.storeLimitSnapshots(d)
if e != nil {
return e
}
rt.dash = d
_ = a.store.UpdateAccount(id, d.Account.Email, d.Account.PlanType, d.Account.Connected)
if resetDetected {
go a.processReminders()
}
return nil
}
const resetDropTolerance = 0.01
func (a *App) storeLimitSnapshots(d Dashboard) (bool, error) {
tx, err := a.store.DB.Begin()
if err != nil {
return false, err
}
defer tx.Rollback()
g := a.general()
resetDetected := false
for _, x := range d.Limits {
var previousID, previousFetchedAt, previousResetsAt int64
var previousUsed float64
err = tx.QueryRow(`SELECT id,used_percent,resets_at,fetched_at FROM limit_snapshots
WHERE account_id=? AND limit_id=? AND window_type=? ORDER BY fetched_at DESC,id DESC LIMIT 1`,
d.AccountID, x.LimitID, x.WindowType).Scan(&previousID, &previousUsed, &previousResetsAt, &previousFetchedAt)
if err != nil && err != sql.ErrNoRows {
return false, err
}
age := d.FetchedAt - previousFetchedAt
if err == nil && g.NotifyAfter && age >= 0 && age <= int64((6*time.Hour).Seconds()) && previousUsed-x.UsedPercent > resetDropTolerance {
kind := "detected_after"
key := fmt.Sprintf("%d:%s:%s:detected:%d", d.AccountID, x.LimitID, x.WindowType, previousID)
now := time.Unix(d.FetchedAt, 0)
if previousResetsAt <= d.FetchedAt && now.Sub(time.Unix(previousResetsAt, 0)) <= 6*time.Hour {
kind = "after"
key = fmt.Sprintf("%d:%s:%s:%d:after", d.AccountID, x.LimitID, x.WindowType, previousResetsAt)
}
body := fmt.Sprintf("Codex [%s] %s/%s 额度已重置:已用 %.1f%% → %.1f%%,下次重置时间 %s。",
d.DisplayName, x.LimitID, x.WindowType, previousUsed, x.UsedPercent, time.Unix(x.ResetsAt, 0).Format(time.RFC3339))
_, err = tx.Exec(`INSERT OR IGNORE INTO notifications
(dedupe_key,channel,kind,status,attempts,last_error,scheduled_at,sent_at,body)
VALUES(?,?,?,'pending',0,'',?,NULL,?)`, key, "configured", kind, d.FetchedAt, body)
if err != nil {
return false, err
}
resetDetected = true
}
if _, err = tx.Exec("INSERT INTO limit_snapshots(limit_id,window_type,used_percent,duration_mins,resets_at,fetched_at,account_id) VALUES(?,?,?,?,?,?,?)", x.LimitID, x.WindowType, x.UsedPercent, x.WindowDurationMinutes, x.ResetsAt, d.FetchedAt, d.AccountID); err != nil {
return false, err
}
}
if err = tx.Commit(); err != nil {
return false, err
}
return resetDetected, nil
}
type rawLimit struct {
LimitID string `json:"limitId"`
LimitName *string `json:"limitName"`
+1
View File
@@ -35,6 +35,7 @@ type App struct {
mu sync.RWMutex
runtimes map[int64]*accountRuntime
loginAttempts sync.Map
reminderMu sync.Mutex
}
type accountRuntime struct {
client codexClient
+50 -31
View File
@@ -341,6 +341,8 @@ func num(n *int64) string {
return fmt.Sprintf("%d", *n)
}
func (a *App) processReminders() {
a.reminderMu.Lock()
defer a.reminderMu.Unlock()
g := a.general()
a.mu.RLock()
ds := make([]Dashboard, 0, len(a.runtimes))
@@ -368,41 +370,58 @@ func (a *App) processReminders() {
continue
}
key := fmt.Sprintf("%d:%s:%s:%d:%s", d.AccountID, x.LimitID, x.WindowType, x.ResetsAt, kind)
var exists int
if a.store.DB.QueryRow("SELECT 1 FROM notifications WHERE dedupe_key=? AND status='sent'", key).Scan(&exists) == nil {
continue
}
body := fmt.Sprintf("Codex [%s] %s/%s 剩余 %.1f%%,重置时间 %s。", d.DisplayName, x.LimitID, x.WindowType, 100-x.UsedPercent, time.Unix(x.ResetsAt, 0).Format(time.RFC3339))
ok := true
errs := []string{}
if t, e := a.telegramSecret(); e == nil && t.Enabled && t.ChatID != 0 {
if e = tgSend(t, body); e != nil {
ok = false
errs = append(errs, e.Error())
}
}
if s, e := a.smtpSecret(); e == nil && s.Enabled {
if e = sendSMTP(s, "Codex 用量重置提醒", body); e != nil {
ok = false
errs = append(errs, e.Error())
}
}
status := "sent"
var sent any = time.Now().Unix()
if !ok {
status = "failed"
sent = nil
}
_, _ = a.store.DB.Exec(`INSERT INTO notifications(dedupe_key,channel,kind,status,attempts,last_error,scheduled_at,sent_at)
VALUES(?,?,?,?,?,?,?,?)
ON CONFLICT(dedupe_key) DO UPDATE SET
status=excluded.status,
attempts=notifications.attempts+1,
last_error=excluded.last_error,
sent_at=excluded.sent_at`, key, "configured", kind, status, 1, strings.Join(errs, "; "), at.Unix(), sent)
_, _ = a.store.DB.Exec(`INSERT OR IGNORE INTO notifications
(dedupe_key,channel,kind,status,attempts,last_error,scheduled_at,sent_at,body)
VALUES(?,?,?,'pending',0,'',?,NULL,?)`, key, "configured", kind, at.Unix(), body)
}
}
}
a.sendPendingReminders(now)
}
func (a *App) sendPendingReminders(now time.Time) {
type pendingReminder struct {
key string
body string
}
rows, err := a.store.DB.Query(`SELECT dedupe_key,body FROM notifications
WHERE status!='sent' AND scheduled_at<=? AND scheduled_at>=? ORDER BY scheduled_at`, now.Unix(), now.Add(-6*time.Hour).Unix())
if err != nil {
return
}
pending := []pendingReminder{}
for rows.Next() {
var p pendingReminder
if rows.Scan(&p.key, &p.body) == nil && p.body != "" {
pending = append(pending, p)
}
}
_ = rows.Close()
for _, p := range pending {
ok := true
errs := []string{}
if t, e := a.telegramSecret(); e == nil && t.Enabled && t.ChatID != 0 {
if e = tgSend(t, p.body); e != nil {
ok = false
errs = append(errs, e.Error())
}
}
if s, e := a.smtpSecret(); e == nil && s.Enabled {
if e = sendSMTP(s, "Codex 用量重置提醒", p.body); e != nil {
ok = false
errs = append(errs, e.Error())
}
}
status := "sent"
var sent any = now.Unix()
if !ok {
status = "failed"
sent = nil
}
_, _ = a.store.DB.Exec(`UPDATE notifications SET status=?,attempts=attempts+1,last_error=?,sent_at=? WHERE dedupe_key=?`,
status, strings.Join(errs, "; "), sent, p.key)
}
}
var _ = bufio.ErrInvalidUnreadByte
+112
View File
@@ -0,0 +1,112 @@
package app
import (
"strconv"
"strings"
"testing"
"time"
"codex-helper/internal/store"
)
func newReminderTestApp(t *testing.T) *App {
t.Helper()
s, err := store.Open(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = s.DB.Close() })
return &App{store: s, runtimes: map[int64]*accountRuntime{}}
}
func reminderDashboard(fetchedAt int64, used float64, resetsAt int64) Dashboard {
return Dashboard{
AccountID: 1,
DisplayName: "测试账号",
FetchedAt: fetchedAt,
Limits: []LimitBucket{{
LimitID: "codex",
WindowType: "primary",
UsedPercent: used,
WindowDurationMinutes: 300,
ResetsAt: resetsAt,
}},
}
}
func notificationCount(t *testing.T, a *App) int {
t.Helper()
var count int
if err := a.store.DB.QueryRow("SELECT COUNT(*) FROM notifications").Scan(&count); err != nil {
t.Fatal(err)
}
return count
}
func TestStoreLimitSnapshotsDetectsEarlyReset(t *testing.T) {
a := newReminderTestApp(t)
now := time.Now().Unix()
if detected, err := a.storeLimitSnapshots(reminderDashboard(now, 42, now+3600)); err != nil || detected {
t.Fatalf("initial snapshot: detected=%v err=%v", detected, err)
}
if detected, err := a.storeLimitSnapshots(reminderDashboard(now+60, 3, now+7200)); err != nil || !detected {
t.Fatalf("reset snapshot: detected=%v err=%v", detected, err)
}
var kind, body string
if err := a.store.DB.QueryRow("SELECT kind,body FROM notifications").Scan(&kind, &body); err != nil {
t.Fatal(err)
}
if kind != "detected_after" || !strings.Contains(body, "42.0% → 3.0%") || !strings.Contains(body, "测试账号") {
t.Fatalf("notification kind=%q body=%q", kind, body)
}
}
func TestStoreLimitSnapshotsUsesScheduledAfterDedupeKey(t *testing.T) {
a := newReminderTestApp(t)
now := time.Now().Unix()
resetAt := now + 30
_, _ = a.storeLimitSnapshots(reminderDashboard(now, 70, resetAt))
detected, err := a.storeLimitSnapshots(reminderDashboard(now+60, 0, now+3600))
if err != nil || !detected {
t.Fatalf("detected=%v err=%v", detected, err)
}
var key, kind string
if err = a.store.DB.QueryRow("SELECT dedupe_key,kind FROM notifications").Scan(&key, &kind); err != nil {
t.Fatal(err)
}
exact := "1:codex:primary:" + strconv.FormatInt(resetAt, 10) + ":after"
if key != exact || kind != "after" {
t.Fatalf("key=%q kind=%q, want %q after", key, kind, exact)
}
}
func TestStoreLimitSnapshotsIgnoresNonResetChanges(t *testing.T) {
tests := []struct {
name string
oldUsed float64
newUsed float64
age time.Duration
notifyAfter bool
}{
{name: "increase", oldUsed: 10, newUsed: 20, age: time.Minute, notifyAfter: true},
{name: "tolerance", oldUsed: 10, newUsed: 9.995, age: time.Minute, notifyAfter: true},
{name: "old snapshot", oldUsed: 50, newUsed: 0, age: 7 * time.Hour, notifyAfter: true},
{name: "disabled", oldUsed: 50, newUsed: 0, age: time.Minute, notifyAfter: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
a := newReminderTestApp(t)
g := defaults()
g.NotifyAfter = tt.notifyAfter
if err := a.store.SetJSON("general", g); err != nil {
t.Fatal(err)
}
now := time.Now().Unix()
_, _ = a.storeLimitSnapshots(reminderDashboard(now, tt.oldUsed, now+3600))
detected, err := a.storeLimitSnapshots(reminderDashboard(now+int64(tt.age.Seconds()), tt.newUsed, now+7200))
if err != nil || detected || notificationCount(t, a) != 0 {
t.Fatalf("detected=%v notifications=%d err=%v", detected, notificationCount(t, a), err)
}
})
}
}