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:
1 parent
9900695855
commit
97a59a2dfb
11 files changed
+671
-6
No files matched your search
@@ -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
@@ -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{
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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,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()
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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.",
|
||||
|
||||
Reference in new issue
Block a user