diff --git a/README.md b/README.md index dd7b882..dcaeac9 100644 --- a/README.md +++ b/README.md @@ -141,6 +141,7 @@ dsfree2api -config config.toml -check | `[limits]` | `max_concurrent_per_site` 每站点并发、`rate_per_minute` 每 Key 每分钟限流（0 = 不限） | | `[proxy]` | `url` 主线路、`fallback_urls` 备用线路、`slow_start_seconds` 首事件超时 | | `[proxypool]` | 代理池总开关、健康检查间隔 / 超时 / URL、裸 `ip:port` 默认协议、Xray 路径 / 版本 / 自动下载；`[proxypool.entries.*]` 手动节点（分享链接或端点）、`[proxypool.subscriptions.*]` 订阅自动拉取 | +| `[quota]` | 配额哨兵：`enabled` / `check_seconds` / `warn_ratio`，定期用池内会话读取站点免费额度，低额告警与重置留档 | | `[upstream]` | 超时、配置缓存 TTL、会话自动刷新、`cross_site_failover` 跨站切换、`continue_rounds` 工具调用截断续写轮数（0 = 不续写、直接截断补全） | | `[turnstile]` | **默认 `enabled = false` 且不内置任何求解服务**；`provider` 三选一：`api`（填 `api_url` / `api_key`）、`browser`（填 `browser_path` 指向本机 Chrome/Edge）、`manual`（只用控制台导入的 Cookie），另有 Cookie TTL、重试次数、单次求解超时（`timeout_seconds`）与 Cookie 池（`warm_enabled` / `warm_ratio` / `warm_check_seconds`，后台预热，请求零等待） | | `[sites.*]` | 三个站点的 `base_url` / `ajax_url` / `sitekey` / `language` / `proxies`（站点级出口绑定，按顺序粘性选路） | diff --git a/cmd/dsfree2api/main.go b/cmd/dsfree2api/main.go index 7edcfd7..1a5207a 100644 --- a/cmd/dsfree2api/main.go +++ b/cmd/dsfree2api/main.go @@ -22,6 +22,7 @@ import ( "github.com/nyoungo/dsfree2api/internal/logfile" "github.com/nyoungo/dsfree2api/internal/metrics" "github.com/nyoungo/dsfree2api/internal/proxypool" + "github.com/nyoungo/dsfree2api/internal/quota" "github.com/nyoungo/dsfree2api/internal/turnstile" "github.com/nyoungo/dsfree2api/internal/upstream" ) @@ -91,14 +92,35 @@ func main() { up := upstream.New(cfg, ts, logger, pool) + // Quota watchdog: every site-side quota exhaustion is logged together with + // the tokens spent on that site so far — including ones swallowed by + // cross-site failover, so the cutoff is recorded rather than hidden. + up.SetQuotaObserver(func(site, model, route string, err error) { + met.IncSiteQuota(site) + t := met.SiteTotals(site) + logger.Warn("site quota exhausted", + "site", site, "model", model, "route", route, "error", err.Error(), + "site_requests", t.Requests, + "site_prompt_tokens", t.PromptTokens, + "site_completion_tokens", t.CompletionTokens, + "site_quota_events", t.Quota) + }) + // Cookie pool: keep one verified session per site × route ahead of time, // so API requests never wait for a solve. No-op unless warm_enabled. warmCtx, warmCancel := context.WithCancel(context.Background()) defer warmCancel() ts.StartWarmer(warmCtx, func() []turnstile.WarmTarget { return warmTargets(cfg, up, pool) }) + // Quota sentinel: poll the sites' guest-token balances for the console + // and the "quota low" alarms. No-op unless [quota].enabled. + quotaCtx, quotaCancel := context.WithCancel(context.Background()) + defer quotaCancel() + quotaWatch := quota.New(cfg, logger, ts, pool) + quotaWatch.Start(quotaCtx) + apiSrv := api.New(cfg, up, met, logs, ts, logger) - adminSrv := admin.New(cfg, up, met, logs, ts, logger, pool) + adminSrv := admin.New(cfg, up, met, logs, ts, logger, pool, quotaWatch) mainAddr := fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port) mainHTTP := &http.Server{ diff --git a/config.example.toml b/config.example.toml index f47d005..5b0ebfd 100644 --- a/config.example.toml +++ b/config.example.toml @@ -117,6 +117,14 @@ warm_enabled = false warm_ratio = 0.3 # 剩余 TTL 低于该比例（0.05–0.95）时触发预热 warm_check_seconds = 60 # 后台巡检间隔秒数（5–3600） +# ── 配额哨兵（站点免费额度实时监控）────────────── +# 用与聊天相同的出口 Cookie 定期查询站点余额接口（每日免费额度），结果 +# 写入日志与后台 /api/overview 的 quota 区域；剩余低于 warn_ratio 时告警。 +[quota] +enabled = false +check_seconds = 300 +warn_ratio = 0.2 + [runtime] data_dir = "./data" diff --git a/internal/admin/server.go b/internal/admin/server.go index e88bf80..fc4dbb3 100644 --- a/internal/admin/server.go +++ b/internal/admin/server.go @@ -16,6 +16,7 @@ import ( "github.com/nyoungo/dsfree2api/internal/logbuf" "github.com/nyoungo/dsfree2api/internal/metrics" "github.com/nyoungo/dsfree2api/internal/proxypool" + "github.com/nyoungo/dsfree2api/internal/quota" "github.com/nyoungo/dsfree2api/internal/turnstile" "github.com/nyoungo/dsfree2api/internal/upstream" ) @@ -33,6 +34,7 @@ type Server struct { logs *logbuf.Buffer ts *turnstile.Solver pool *proxypool.Manager + quota *quota.Watcher log *slog.Logger start time.Time @@ -41,7 +43,7 @@ type Server struct { password string } -func New(cfg *config.Config, up *upstream.Client, met *metrics.Recorder, logs *logbuf.Buffer, ts *turnstile.Solver, log *slog.Logger, pool *proxypool.Manager) *Server { +func New(cfg *config.Config, up *upstream.Client, met *metrics.Recorder, logs *logbuf.Buffer, ts *turnstile.Solver, log *slog.Logger, pool *proxypool.Manager, qw *quota.Watcher) *Server { if log == nil { log = slog.Default() } @@ -52,7 +54,7 @@ func New(cfg *config.Config, up *upstream.Client, met *metrics.Recorder, logs *l generated = true } s := &Server{ - cfg: cfg, up: up, met: met, logs: logs, ts: ts, pool: pool, log: log, + cfg: cfg, up: up, met: met, logs: logs, ts: ts, pool: pool, quota: qw, log: log, start: time.Now(), tokens: map[string]time.Time{}, password: pw, } if generated { @@ -207,6 +209,11 @@ func (s *Server) handleOverview(w http.ResponseWriter, _ *http.Request) { if s.pool != nil { poolStatus = s.pool.Status() } + var quotaStatus any + if s.quota != nil { + quotaStatus = s.quota.Snapshot() + } + writeJSON(w, http.StatusOK, map[string]any{ "metrics": s.met.Snapshot(), "models": models, @@ -214,6 +221,7 @@ func (s *Server) handleOverview(w http.ResponseWriter, _ *http.Request) { "proxy": proxyCfg, "turnstile": map[string]any{"config": turnstileCfg, "status": s.ts.Status(siteCodes), "history": s.ts.History(), "pool": s.ts.PoolStatus()}, "proxypool": poolStatus, + "quota": quotaStatus, "keys": keys, "cache": map[string]any{"chat_config_entries": s.up.ConfigCacheSize(), "routes": routeList(s.up)}, "system": map[string]any{ diff --git a/internal/config/config.go b/internal/config/config.go index 394a2a8..f5fc313 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -126,6 +126,53 @@ func (t Turnstile) WarmCheckValue() int { return n } +// Quota sentinel defaults: polling interval and the remaining-quota warning +// threshold (share of the daily free limit left). +const ( + DefaultQuotaCheckSeconds = 300 + DefaultQuotaWarnRatio = 0.2 +) + +// QuotaWatch configures the guest-token balance sentinel: it polls each +// site's balance endpoint with the pooled session (same egress and cookies as +// chat) and reports remaining daily quota to the console and logs. +type QuotaWatch struct { + Enabled bool `toml:"enabled" json:"enabled"` + CheckSeconds int `toml:"check_seconds" json:"check_seconds"` + WarnRatio float64 `toml:"warn_ratio" json:"warn_ratio"` +} + +// CheckSecondsValue normalizes check_seconds (default 300s, clamped 30–3600). +func (q QuotaWatch) CheckSecondsValue() int { + n := q.CheckSeconds + if n <= 0 { + return DefaultQuotaCheckSeconds + } + if n < 30 { + return 30 + } + if n > 3600 { + return 3600 + } + return n +} + +// WarnRatioValue normalizes warn_ratio: remaining/limit below this fraction +// raises a "quota low" warning (default 20%, clamped 5–95%). +func (q QuotaWatch) WarnRatioValue() float64 { + r := q.WarnRatio + if r <= 0 { + return DefaultQuotaWarnRatio + } + if r < 0.05 { + return 0.05 + } + if r > 0.95 { + return 0.95 + } + return r +} + type Site struct { Code string `toml:"-"` BaseURL string `toml:"base_url" json:"base_url"` @@ -214,6 +261,11 @@ type Config struct { } `toml:"proxy" json:"proxy"` ProxyPool ProxyPool `toml:"proxypool" json:"proxypool"` + + // Quota is the balance sentinel: it polls each site's guest-token balance + // so the console shows how much daily free quota is left. + Quota QuotaWatch `toml:"quota" json:"quota"` + Upstream struct { Timeout float64 `toml:"timeout" json:"timeout"` StreamTimeout float64 `toml:"stream_timeout" json:"stream_timeout"` diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index a3d8d00..e23acc3 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -46,9 +46,11 @@ type modelStats struct { } type siteStats struct { - Requests int64 `json:"requests"` - Errors int64 `json:"errors"` - Quota int64 `json:"quota_exhausted"` + Requests int64 `json:"requests"` + Errors int64 `json:"errors"` + Quota int64 `json:"quota_exhausted"` + Prompt int64 `json:"prompt_tokens"` + Completed int64 `json:"completion_tokens"` } type dayStats struct { @@ -228,6 +230,8 @@ func (r *Recorder) Record(rec Record) { r.sites[rec.Site] = s } s.Requests++ + s.Prompt += int64(rec.PromptTokens) + s.Completed += int64(rec.CompletionTokens) if rec.Status == StatusError || rec.Status == StatusEmpty { s.Errors++ } @@ -264,6 +268,45 @@ func (r *Recorder) Record(rec Record) { } } +// SiteTotals is the cumulative per-site usage snapshot; it sizes a quota +// exhaustion — how many tokens had flowed before the site cut us off. +type SiteTotals struct { + Requests int64 `json:"requests"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + Quota int64 `json:"quota_exhausted"` +} + +// SiteTotals returns the cumulative usage attributed to one site so far. +// Totals persist across restarts through metrics.json. +func (r *Recorder) SiteTotals(site string) SiteTotals { + r.mu.Lock() + defer r.mu.Unlock() + s := r.sites[site] + if s == nil { + return SiteTotals{} + } + return SiteTotals{ + Requests: s.Requests, + PromptTokens: s.Prompt, + CompletionTokens: s.Completed, + Quota: s.Quota, + } +} + +// IncSiteQuota counts one site-side quota exhaustion that may never surface as +// a failed request (e.g. swallowed by cross-site failover). +func (r *Recorder) IncSiteQuota(site string) { + r.mu.Lock() + defer r.mu.Unlock() + s := r.sites[site] + if s == nil { + s = &siteStats{} + r.sites[site] = s + } + s.Quota++ +} + func (r *Recorder) Snapshot() any { r.mu.Lock() defer r.mu.Unlock() diff --git a/internal/metrics/metrics_test.go b/internal/metrics/metrics_test.go new file mode 100644 index 0000000..04e9553 --- /dev/null +++ b/internal/metrics/metrics_test.go @@ -0,0 +1,26 @@ +package metrics + +import "testing" + +func TestRecordTracksSiteTokensAndQuota(t *testing.T) { + r := New(t.TempDir()) + r.Record(Record{Model: "deepseek-v4-flash-de", Site: "de", PromptTokens: 120, CompletionTokens: 80, Status: StatusOK}) + r.Record(Record{Model: "deepseek-v4-flash-de", Site: "de", PromptTokens: 15, Status: StatusError, Error: "sse quota exhausted"}) + + got := r.SiteTotals("de") + if got.Requests != 2 { + t.Fatalf("Requests = %d, want 2", got.Requests) + } + if got.PromptTokens != 135 || got.CompletionTokens != 80 { + t.Errorf("tokens = %d/%d, want 135/80", got.PromptTokens, got.CompletionTokens) + } + + r.IncSiteQuota("de") + if got := r.SiteTotals("de"); got.Quota != 1 { + t.Errorf("Quota = %d, want 1", got.Quota) + } + + if empty := r.SiteTotals("fr"); empty.Requests != 0 || empty.Quota != 0 { + t.Errorf("unknown site totals = %+v, want zero", empty) + } +} diff --git a/internal/quota/watcher.go b/internal/quota/watcher.go new file mode 100644 index 0000000..d315db1 --- /dev/null +++ b/internal/quota/watcher.go @@ -0,0 +1,357 @@ +// Package quota watches the upstream sites' guest-token balances so the +// console and logs show how much free daily quota is left before a site cuts +// requests off. Readings reuse the pooled cookies and egress of the chat +// path, so they describe the identity requests are actually debited from. +package quota + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "sort" + "strconv" + "strings" + "sync" + "time" + + "github.com/nyoungo/dsfree2api/internal/config" + "github.com/nyoungo/dsfree2api/internal/httpx" + "github.com/nyoungo/dsfree2api/internal/proxypool" + "github.com/nyoungo/dsfree2api/internal/turnstile" +) + +// Balance is one site × bot quota reading. Empty Error means the reading was +// taken successfully; Limit/Used/Remaining come from the site's own balance +// endpoint (the guest daily free allowance). +type Balance struct { + Site string `json:"site"` + BotID int `json:"bot_id"` + Model string `json:"model,omitempty"` + Route string `json:"route"` + PaidBalance int64 `json:"paid_balance"` + Limit int64 `json:"limit"` + Used int64 `json:"used"` + Remaining int64 `json:"remaining"` + Ratio float64 `json:"ratio"` + ResetPeriod string `json:"reset_period,omitempty"` + CheckedAt int64 `json:"checked_at"` + Error string `json:"error,omitempty"` +} + +// Snapshot is the console view of the sentinel. +type Snapshot struct { + Enabled bool `json:"enabled"` + Running bool `json:"running"` + CheckS int `json:"check_seconds"` + WarnRatio float64 `json:"warn_ratio"` + Entries []Balance `json:"entries"` +} + +type botRef struct { + id int + model string +} + +// Watcher polls each site's balance REST endpoint on an interval. One slot +// per site × bot; slots without verified cookies are reported as such and +// never trigger solves by themselves. +type Watcher struct { + cfg *config.Config + log *slog.Logger + ts *turnstile.Solver + pool *proxypool.Manager + + // newSession is a test seam; nil uses the real Chrome-impersonating one. + newSession func(proxy string) (httpx.Session, error) + + mu sync.Mutex + entries map[string]*Balance // key: site|bot + warned map[string]bool + lastErr map[string]string + running bool +} + +// New builds a watcher. Start launches the loop when [quota].enabled. +func New(cfg *config.Config, log *slog.Logger, ts *turnstile.Solver, pool *proxypool.Manager) *Watcher { + if log == nil { + log = slog.Default() + } + return &Watcher{ + cfg: cfg, + log: log, + ts: ts, + pool: pool, + entries: map[string]*Balance{}, + warned: map[string]bool{}, + lastErr: map[string]string{}, + } +} + +// Start launches the polling loop; no-op unless [quota].enabled. +func (w *Watcher) Start(ctx context.Context) { + w.cfg.RLock() + enabled := w.cfg.Quota.Enabled + w.cfg.RUnlock() + if !enabled || w.ts == nil { + return + } + go w.loop(ctx) +} + +func (w *Watcher) loop(ctx context.Context) { + w.mu.Lock() + w.running = true + w.mu.Unlock() + defer func() { + w.mu.Lock() + w.running = false + w.mu.Unlock() + }() + // First sweep after a short settle delay (pool/cookies may still be + // warming up), then on every check interval. + timer := time.NewTimer(20 * time.Second) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + w.Sweep(ctx) + w.cfg.RLock() + iv := time.Duration(w.cfg.Quota.CheckSecondsValue()) * time.Second + w.cfg.RUnlock() + timer.Reset(iv) + } + } +} + +// Sweep queries every enabled site once. Exposed for tests and manual runs. +func (w *Watcher) Sweep(ctx context.Context) { + if w.ts == nil { + return + } + w.cfg.RLock() + sites := map[string]config.Site{} + for code, st := range w.cfg.Sites { + if st != nil && st.Enabled { + sites[code] = *st + } + } + bots := map[string][]botRef{} + for id, m := range w.cfg.Models { + if m == nil || !m.Enabled || m.BotID <= 0 { + continue + } + bots[m.Site] = append(bots[m.Site], botRef{id: m.BotID, model: id}) + } + warnRatio := w.cfg.Quota.WarnRatioValue() + w.cfg.RUnlock() + + codes := make([]string, 0, len(sites)) + for code := range sites { + codes = append(codes, code) + } + sort.Strings(codes) + for _, code := range codes { + refs := bots[code] + if len(refs) == 0 { + continue + } + sort.Slice(refs, func(i, j int) bool { return refs[i].id < refs[j].id }) + w.checkSite(ctx, sites[code], refs, warnRatio) + } +} + +func (w *Watcher) checkSite(ctx context.Context, site config.Site, refs []botRef, warnRatio float64) { + route := w.routeFor(site.Code) + if !w.ts.CookiesReady(site.Code, route) { + w.markAll(site.Code, refs, route, errors.New("no verified cookies for this route yet")) + return + } + sess, err := w.session(route) + if err != nil { + w.markAll(site.Code, refs, route, err) + return + } + defer sess.Close() + if err := w.ts.ApplyValidCookies(ctx, sess, &site, route); err != nil { + w.markAll(site.Code, refs, route, err) + return + } + // Guests need no REST nonce: the visitor id travels as the dsgt_gid + // cookie and the balance endpoint keys on it. + if gid := sess.Cookies()["dsgt_gid"]; gid == "" { + w.markAll(site.Code, refs, route, errors.New("no dsgt_gid cookie for this route yet")) + return + } + endpoint := strings.TrimRight(site.BaseURL, "/") + "/wp-json/dsgt/v1/balance" + for _, ref := range refs { + payload, err := fetchBalance(ctx, sess, endpoint, ref.id) + w.record(site.Code, ref, route, payload, err, warnRatio) + } +} + +func (w *Watcher) markAll(site string, refs []botRef, route string, err error) { + for _, ref := range refs { + w.record(site, ref, route, nil, err, 0) + } +} + +func (w *Watcher) record(site string, ref botRef, route string, payload *balancePayload, err error, warnRatio float64) { + key := site + "|" + strconv.Itoa(ref.id) + b := &Balance{ + Site: site, BotID: ref.id, Model: ref.model, + Route: routeLabel(route), CheckedAt: time.Now().Unix(), + } + if err != nil { + b.Error = err.Error() + } else { + b.PaidBalance = payload.Balance + b.Limit = payload.Free.Limit + b.Used = payload.Free.Used + b.Remaining = payload.Free.Remaining + b.ResetPeriod = payload.Free.ResetPeriod + if b.Limit > 0 { + b.Ratio = float64(b.Remaining) / float64(b.Limit) + } + } + + var lowNow, recovered, resetSeen bool + var prevUsed int64 + w.mu.Lock() + if prev := w.entries[key]; prev != nil && prev.Error == "" && err == nil && b.Used < prev.Used { + resetSeen = true + prevUsed = prev.Used + } + w.entries[key] = b + prevErr := w.lastErr[key] + w.lastErr[key] = b.Error + if err == nil && b.Limit > 0 { + low := float64(b.Remaining) <= warnRatio*float64(b.Limit) + if low && !w.warned[key] { + w.warned[key] = true + lowNow = true + } + if !low && w.warned[key] { + w.warned[key] = false + recovered = true + } + } + w.mu.Unlock() + + if err != nil { + if b.Error != prevErr { + w.log.Warn("quota check failed", + "site", site, "bot_id", ref.id, "model", ref.model, "route", b.Route, "error", b.Error) + } + return + } + if prevErr != "" { + w.log.Info("quota check recovered", "site", site, "bot_id", ref.id, "route", b.Route) + } + w.log.Debug("quota check", + "site", site, "bot_id", ref.id, "model", ref.model, "route", b.Route, + "remaining", b.Remaining, "limit", b.Limit, "used", b.Used) + if resetSeen { + w.log.Info("quota reset observed", + "site", site, "bot_id", ref.id, "model", ref.model, "route", b.Route, + "used_before", prevUsed, "used_after", b.Used, + "checked_at", time.Unix(b.CheckedAt, 0).Format(time.RFC3339)) + } + switch { + case lowNow: + w.log.Warn("quota low", + "site", site, "bot_id", ref.id, "model", ref.model, "route", b.Route, + "remaining", b.Remaining, "limit", b.Limit, "reset_period", b.ResetPeriod) + case recovered: + w.log.Info("quota recovered", + "site", site, "bot_id", ref.id, "model", ref.model, + "remaining", b.Remaining, "limit", b.Limit) + } +} + +func (w *Watcher) routeFor(site string) string { + if w.pool != nil { + if cands := w.pool.Candidates(site); len(cands) > 0 { + return cands[0].Proxy + } + } + return "" +} + +func (w *Watcher) session(proxy string) (httpx.Session, error) { + if w.newSession != nil { + return w.newSession(proxy) + } + return httpx.NewSession(httpx.Options{ProxyURL: proxy, Timeout: 45 * time.Second}) +} + +func routeLabel(proxy string) string { + if proxy == "" { + return "direct" + } + return proxy +} + +// freeQuota is the daily guest allowance block of a balance response. +type freeQuota struct { + Limit int64 `json:"limit"` + Used int64 `json:"used"` + Remaining int64 `json:"remaining"` + ResetPeriod string `json:"reset_period"` +} + +type balancePayload struct { + Balance int64 `json:"balance"` + Free *freeQuota `json:"free"` +} + +func fetchBalance(ctx context.Context, sess httpx.Session, endpoint string, botID int) (*balancePayload, error) { + u := fmt.Sprintf("%s?bot_id=%d", endpoint, botID) + resp, err := sess.Do(ctx, httpx.Request{ + Method: http.MethodGet, + URL: u, + Header: map[string]string{"Accept": "application/json"}, + }) + if err != nil { + return nil, fmt.Errorf("balance request: %w", err) + } + if resp.Status != http.StatusOK { + return nil, fmt.Errorf("balance request: HTTP %d", resp.Status) + } + var p balancePayload + if err := json.Unmarshal(resp.Body, &p); err != nil { + return nil, fmt.Errorf("balance json: %w", err) + } + if p.Free == nil { + return nil, errors.New("balance response has no free quota block") + } + return &p, nil +} + +// Snapshot returns the current readings for the console. +func (w *Watcher) Snapshot() Snapshot { + w.cfg.RLock() + enabled := w.cfg.Quota.Enabled + checkS := w.cfg.Quota.CheckSecondsValue() + ratio := w.cfg.Quota.WarnRatioValue() + w.cfg.RUnlock() + + w.mu.Lock() + entries := make([]Balance, 0, len(w.entries)) + for _, b := range w.entries { + entries = append(entries, *b) + } + running := w.running + w.mu.Unlock() + sort.Slice(entries, func(i, j int) bool { + if entries[i].Site != entries[j].Site { + return entries[i].Site < entries[j].Site + } + return entries[i].BotID < entries[j].BotID + }) + return Snapshot{Enabled: enabled, Running: running, CheckS: checkS, WarnRatio: ratio, Entries: entries} +} diff --git a/internal/quota/watcher_test.go b/internal/quota/watcher_test.go new file mode 100644 index 0000000..039cdd0 --- /dev/null +++ b/internal/quota/watcher_test.go @@ -0,0 +1,105 @@ +package quota + +import ( + "bytes" + "context" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/nyoungo/dsfree2api/internal/config" + "github.com/nyoungo/dsfree2api/internal/httpx" + "github.com/nyoungo/dsfree2api/internal/turnstile" +) + +const examplePath = "../../config.example.toml" + +func TestQuotaResetIsRecorded(t *testing.T) { + cfg, err := config.Load(examplePath) + if err != nil { + t.Fatalf("load config: %v", err) + } + var buf bytes.Buffer + w := New(cfg, slog.New(slog.NewTextHandler(&buf, nil)), nil, nil) + ref := botRef{id: 7, model: "m"} + w.record("de", ref, "", &balancePayload{Free: &freeQuota{Limit: 30000, Used: 500, Remaining: 29500, ResetPeriod: "daily"}}, nil, 0.2) + w.record("de", ref, "", &balancePayload{Free: &freeQuota{Limit: 30000, Used: 0, Remaining: 30000, ResetPeriod: "daily"}}, nil, 0.2) + if !strings.Contains(buf.String(), "quota reset observed") { + t.Fatalf("expected a reset log line, got:\n%s", buf.String()) + } + if !strings.Contains(buf.String(), "used_before=500") { + t.Fatalf("expected used_before=500 in log, got:\n%s", buf.String()) + } +} + +func TestSweepReadsBalancesWithPooledCookies(t *testing.T) { + var seenGid string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasPrefix(r.URL.Path, "/wp-json/dsgt/v1/balance") { + t.Errorf("unexpected request %s — the sentinel must not fetch the site page", r.URL.Path) + return + } + if c, err := r.Cookie("dsgt_gid"); err == nil { + seenGid = c.Value + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"balance":0,"free":{"limit":30000,"used":500,"remaining":29500,"reset_period":"daily"}}`)) + })) + defer srv.Close() + + cfg, err := config.Load(examplePath) + if err != nil { + t.Fatalf("load config: %v", err) + } + cfg.Quota.Enabled = true + cfg.Quota.CheckSeconds = 300 + cfg.Quota.WarnRatio = 0.2 + site := *cfg.Sites["de"] + site.BaseURL = srv.URL + cfg.Sites["de"] = &site + + ts := turnstile.New(cfg.Turnstile, "") + if _, _, err := ts.ImportCookies("de", "dsgt_gid=test-gid; cf_clearance=abc", ""); err != nil { + t.Fatalf("import cookies: %v", err) + } + + w := New(cfg, nil, ts, nil) + w.newSession = func(string) (httpx.Session, error) { + return httpx.NewSession(httpx.Options{Timeout: 5 * time.Second}) + } + w.Sweep(context.Background()) + + if seenGid != "test-gid" { + t.Fatalf("balance call visitor cookie = %q, want test-gid", seenGid) + } + snap := w.Snapshot() + if !snap.Enabled { + t.Error("snapshot should report enabled") + } + if len(snap.Entries) != 6 { + t.Fatalf("entries = %d, want 6 (3 sites × 2 bots)", len(snap.Entries)) + } + var de *Balance + for i := range snap.Entries { + if snap.Entries[i].Site == "de" && snap.Entries[i].BotID == 27487 { + de = &snap.Entries[i] + } + } + if de == nil { + t.Fatalf("no de/27487 entry in %+v", snap.Entries) + } + if de.Error != "" || de.Limit != 30000 || de.Remaining != 29500 || de.Used != 500 || de.ResetPeriod != "daily" { + t.Fatalf("de balance = %+v", de) + } + if de.Route != "direct" { + t.Errorf("route = %q, want direct", de.Route) + } + for _, e := range snap.Entries { + if e.Site != "de" && e.Error == "" { + t.Errorf("site %s should report missing cookies, got %+v", e.Site, e) + } + } +} diff --git a/internal/upstream/client.go b/internal/upstream/client.go index ff62ce5..f620a1d 100644 --- a/internal/upstream/client.go +++ b/internal/upstream/client.go @@ -73,6 +73,11 @@ type Client struct { lastWarnMsg string lastWarnAt time.Time + // quotaObserver, when set, is notified on the first site-side quota + // exhaustion of each request — including ones hidden by cross-site + // failover. + quotaObserver func(site, model, route string, err error) + // chatOnceOverride is a test seam; nil in production. chatOnceOverride func(ctx context.Context, site config.Site, modelID string, model config.Model, prompt string, route Route, info *ServeInfo, yield func(Event) error) error @@ -94,6 +99,13 @@ func New(cfg *config.Config, ts *turnstile.Solver, log *slog.Logger, pool RouteS } } +// SetQuotaObserver registers a callback that fires whenever a site reports +// quota exhaustion, even when cross-site failover later succeeds and the error +// would otherwise never surface. Set it before serving requests. +func (c *Client) SetQuotaObserver(fn func(site, model, route string, err error)) { + c.quotaObserver = fn +} + // Routes returns the primary proxy followed by the fallbacks, deduplicated. func (c *Client) Routes() []Route { c.cfg.RLock() @@ -405,6 +417,7 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf return errf("site %s is disabled", model.Site) } routes := c.routesFor(site.Code) + quotaNotified := false var lastErr error sawQuota := false @@ -426,6 +439,10 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf return nil } lastErr = err + if isQuota(err) && !quotaNotified && c.quotaObserver != nil { + quotaNotified = true + c.quotaObserver(site.Code, modelID, route.Name, err) + } if contentStarted(err) { return AsUpstream(err) } diff --git a/internal/upstream/client_test.go b/internal/upstream/client_test.go index 8e1d3df..45f37d3 100644 --- a/internal/upstream/client_test.go +++ b/internal/upstream/client_test.go @@ -144,6 +144,32 @@ func TestNoFallbackAfterContentStarted(t *testing.T) { } } +func TestQuotaObserverFiresOncePerSiteBeforeFailover(t *testing.T) { + cfg := testConfig(t) + cfg.Upstream.AutoRefresh = false + c := New(cfg, nil, nil, nil) + c.chatOnceOverride = func(ctx context.Context, site config.Site, modelID string, model config.Model, + prompt string, route Route, info *ServeInfo, yield func(Event) error) error { + if site.Code == "de" { + return newQuota(`sse quota exhausted: {"quota_notice":{}}`) + } + yield(Event{Kind: KindDelta, Value: "ok"}) + return nil + } + var fired []string + c.SetQuotaObserver(func(site, model, route string, err error) { + fired = append(fired, site+"|"+model+"|"+route) + }) + + if err := c.Chat(context.Background(), "deepseek-v4-flash-de", "hi", &ServeInfo{}, + func(ev Event) error { return nil }); err != nil { + t.Fatalf("chat should succeed via cross-site failover: %v", err) + } + if len(fired) != 1 || fired[0] != "de|deepseek-v4-flash-de|primary" { + t.Fatalf("quota observer events = %v, want [de|deepseek-v4-flash-de|primary]", fired) + } +} + func TestQuotaNoticeTriggersQuotaError(t *testing.T) { payload := map[string]any{ "error": "Dein heutiges Guthaben von 10.000 Token (V4-Pro) ist aufgebraucht.",