feat(quota): 配额哨兵与配额耗尽留档

- 新增 internal/quota 哨兵:[quota] 定期用池内身份查询站点余额(访客态无需 nonce),
  quota low 告警、quota reset observed 留档,/api/overview 暴露 site × bot 余额
- upstream: SetQuotaObserver 观测器(含被跨站接力吞没的配额事件)
- metrics: 按站点累计 prompt / completion tokens 与配额事件计数
This commit is contained in:
luotianyi committed 2026-10-08 00:19:40 +08:00
1 parent 9900695855
commit 97a59a2dfb
11 files changed
+671 -6

No files matched your search

+1
View File
@@ -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`(站点级出口绑定,按顺序粘性选路) |
+23 -1
View File
@@ -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{
+8
View File
@@ -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"
+10 -2
View File
@@ -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{
+52
View File
@@ -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"`
+46 -3
View File
@@ -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()
+26
View File
@@ -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)
}
}
+357
View File
@@ -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}
}
+105
View File
@@ -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)
}
}
}
+17
View File
@@ -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)
}
+26
View File
@@ -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.",