feat: rotate visitor identity to restore a quota-exhausted site without re-solving

internal/turnstile/solver.go

@@ -5,6 +5,7 @@ package turnstile
 import (
 	"bytes"
 	"context"
+	"crypto/rand"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -158,6 +159,64 @@ func (s *Solver) Invalidate(siteCode string) {
 	}
 }
 
+// guestCookieName is the cookie the upstream guest-token plugin keys a
+// visitor's free daily tier to (verified against the live endpoints: the
+// balance API echoes it as "gid" and accepts any freshly minted value).
+const guestCookieName = "dsgt_gid"
+
+// guestIDAlphabet matches the 32-char ids the site itself issues.
+const guestIDAlphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+// randomGuestID mints a visitor id in the same shape the site uses.
+func randomGuestID() string {
+	b := make([]byte, 32)
+	if _, err := rand.Read(b); err != nil {
+		return fmt.Sprintf("g%d", time.Now().UnixNano())
+	}
+	for i := range b {
+		b[i] = guestIDAlphabet[int(b[i])%len(guestIDAlphabet)]
+	}
+	return string(b)
+}
+
+// RotateGuest swaps the visitor cookie inside every cached session for the
+// site. The upstream grants each fresh visitor id its own daily free tier —
+// the same effect as opening a new private window — so a quota-exhausted
+// site recovers without paying for a Turnstile re-solve. Returns the new id,
+// or "" when the site has no live cached session to rotate.
+func (s *Solver) RotateGuest(siteCode string) string {
+	gid := randomGuestID()
+	s.mu.Lock()
+	cores := make([]*core, 0, len(s.cores))
+	for _, c := range s.cores {
+		cores = append(cores, c)
+	}
+	s.mu.Unlock()
+	rotated := false
+	for _, c := range cores {
+		c.mu.Lock()
+		if st, ok := c.cache[siteCode]; ok && time.Now().Before(st.expiresAt) {
+			st.cookies[guestCookieName] = gid
+			rotated = true
+		}
+		c.mu.Unlock()
+	}
+	if !rotated {
+		return ""
+	}
+	return gid
+}
+
+// GuestID reports the visitor id cached for the site on one proxy route.
+func (s *Solver) GuestID(siteCode, proxy string) (string, bool) {
+	ck := s.coreFor(proxy).load(siteCode)
+	if ck == nil {
+		return "", false
+	}
+	gid := ck[guestCookieName]
+	return gid, gid != ""
+}
+
 // ImportCookies stores cookies exported from a browser session that already
 // passed the Turnstile check, so chat works while [turnstile] is off or the
 // solver service is down. Cookies are bound to the IP/UA that obtained them,
@@ -211,8 +270,8 @@ func ParseCookiePairs(raw string) map[string]string {
 // [turnstile] switch is off; only the solving step needs the switch.
 func (s *Solver) ApplyValidCookies(ctx context.Context, sess httpx.Session, site *config.Site, proxy string) error {
 	c := s.coreFor(proxy)
-	if st := c.load(site.Code); st != nil {
-		sess.SetCookies(st.cookies)
+	if ck := c.load(site.Code); ck != nil {
+		sess.SetCookies(ck)
 		return nil
 	}
 	if !s.Enabled() {
@@ -221,8 +280,8 @@ func (s *Solver) ApplyValidCookies(ctx context.Context, sess httpx.Session, site
 	l := c.lockFor(site.Code)
 	l.Lock()
 	defer l.Unlock()
-	if st := c.load(site.Code); st != nil {
-		sess.SetCookies(st.cookies)
+	if ck := c.load(site.Code); ck != nil {
+		sess.SetCookies(ck)
 		return nil
 	}
 	return s.refreshLocked(ctx, sess, site, c)
@@ -240,14 +299,20 @@ func (s *Solver) Refresh(ctx context.Context, sess httpx.Session, site *config.S
 	return s.refreshLocked(ctx, sess, site, c)
 }
 
-func (c *core) load(site string) *cookieState {
+// load returns a copy of the cached cookies, or nil when absent or expired.
+// Callers get a snapshot, so RotateGuest can swap entries concurrently.
+func (c *core) load(site string) map[string]string {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	st, ok := c.cache[site]
 	if !ok || time.Now().After(st.expiresAt) {
 		return nil
 	}
-	return st
+	out := make(map[string]string, len(st.cookies))
+	for k, v := range st.cookies {
+		out[k] = v
+	}
+	return out
 }
 
 func (s *Solver) refreshLocked(ctx context.Context, sess httpx.Session, site *config.Site, c *core) error {

internal/turnstile/solver_test.go

@@ -192,3 +192,50 @@ func TestHistoryIsCappedAndRecorded(t *testing.T) {
 		t.Error("last record should be OK")
 	}
 }
+
+func TestRotateGuestSwapsVisitorID(t *testing.T) {
+	cfg, err := config.Load(examplePath)
+	if err != nil {
+		t.Fatalf("load config: %v", err)
+	}
+	solver := New(cfg.Turnstile, httpx.DefaultUserAgent)
+	const oldGid = "OLDGIDOLDGIDOLDGIDOLDGIDOLDG12"
+	if _, _, err := solver.ImportCookies("de", "cf_clearance=abc; dsgt_gid="+oldGid, ""); err != nil {
+		t.Fatalf("import: %v", err)
+	}
+	gid, ok := solver.GuestID("de", "")
+	if !ok || gid != oldGid {
+		t.Fatalf("GuestID = %q %v, want %q", gid, ok, oldGid)
+	}
+	sess := &fakeSession{}
+	if err := solver.ApplyValidCookies(context.Background(), sess, cfg.Sites["de"], ""); err != nil {
+		t.Fatalf("apply: %v", err)
+	}
+	if sess.Cookies()["dsgt_gid"] != oldGid {
+		t.Fatalf("session gid = %q, want the cached %q", sess.Cookies()["dsgt_gid"], oldGid)
+	}
+
+	rotated := solver.RotateGuest("de")
+	if rotated == "" {
+		t.Fatal("RotateGuest returned empty for a live session")
+	}
+	if rotated == oldGid {
+		t.Fatal("RotateGuest kept the old id")
+	}
+	if len(rotated) != 32 {
+		t.Errorf("rotated id length = %d, want 32", len(rotated))
+	}
+	if got, _ := solver.GuestID("de", ""); got != rotated {
+		t.Errorf("GuestID after rotate = %q, want %q", got, rotated)
+	}
+	sess2 := &fakeSession{}
+	if err := solver.ApplyValidCookies(context.Background(), sess2, cfg.Sites["de"], ""); err != nil {
+		t.Fatalf("apply 2: %v", err)
+	}
+	if sess2.Cookies()["dsgt_gid"] != rotated {
+		t.Errorf("fresh session gid = %q, want rotated %q", sess2.Cookies()["dsgt_gid"], rotated)
+	}
+	if got := solver.RotateGuest("no-such-site"); got != "" {
+		t.Errorf("rotating an unknown site returned %q, want empty", got)
+	}
+}

internal/upstream/balance.go

@@ -0,0 +1,78 @@
+package upstream
+
+import (
+	"context"
+	"encoding/json"
+	"strconv"
+	"time"
+
+	"github.com/nyoungo/dsfree2api/internal/config"
+	"github.com/nyoungo/dsfree2api/internal/httpx"
+)
+
+// guestCookieName mirrors the upstream guest-token plugin's visitor cookie;
+// its value is the "gid" the balance endpoint reports on.
+const guestCookieName = "dsgt_gid"
+
+// SiteBalance is the response of GET /wp-json/dsgt/v1/balance. The free-tier
+// block only appears when the request carries a bot_id.
+type SiteBalance struct {
+	Gid     string  `json:"gid"`
+	Balance float64 `json:"balance"`
+	Free    struct {
+		Limit       int    `json:"limit"`
+		Used        int    `json:"used"`
+		Remaining   int    `json:"remaining"`
+		ResetPeriod string `json:"reset_period"`
+	} `json:"free"`
+}
+
+// FetchBalance queries the site's guest balance endpoint for one visitor id.
+// No nonce is required for guests; the id travels as the visitor cookie.
+// botID selects the per-bot free tier (0 omits it, which also omits free).
+func (c *Client) FetchBalance(
+	ctx context.Context,
+	site config.Site,
+	gid string,
+	botID int,
+	route Route,
+) (*SiteBalance, error) {
+	if gid == "" {
+		return nil, errf("no guest id for site %s", site.Code)
+	}
+	sess, err := httpx.NewSession(httpx.Options{ProxyURL: route.Proxy, Timeout: 20 * time.Second})
+	if err != nil {
+		return nil, err
+	}
+	defer sess.Close()
+	sess.SetCookies(map[string]string{guestCookieName: gid})
+
+	url := site.BaseURL + "/wp-json/dsgt/v1/balance?_t=" + strconv.FormatInt(time.Now().UnixMilli(), 10)
+	if botID > 0 {
+		url += "&bot_id=" + itoa(botID)
+	}
+	c.cfg.RLock()
+	ua := c.cfg.Upstream.UserAgent
+	c.cfg.RUnlock()
+	resp, err := sess.Do(ctx, httpx.Request{
+		Method: "GET",
+		URL:    url,
+		Header: map[string]string{
+			"User-Agent": ua,
+			"Accept":     "application/json",
+			"Referer":    site.BaseURL + "/",
+		},
+		Timeout: 20 * time.Second,
+	})
+	if err != nil {
+		return nil, err
+	}
+	if resp.Status >= 400 {
+		return nil, errf("balance http %d: %s", resp.Status, truncate(string(resp.Body), 200))
+	}
+	var out SiteBalance
+	if err := json.Unmarshal(resp.Body, &out); err != nil {
+		return nil, errf("balance bad json: %v", err)
+	}
+	return &out, nil
+}

internal/upstream/client.go

@@ -156,6 +156,43 @@ func (c *Client) sortQuotaCoolLast(ids []string) []string {
 	return append(rest, cooled...)
 }
 
+// clearQuotaCool lifts a site's cooldown once it serves again after a quota
+// error — the fresh visitor identity restored its daily tier.
+func (c *Client) clearQuotaCool(code string) {
+	c.stateMu.Lock()
+	_, ok := c.quotaCool[code]
+	if ok {
+		delete(c.quotaCool, code)
+	}
+	c.stateMu.Unlock()
+	if ok {
+		c.log.Info("site quota restored", "site", code)
+	}
+}
+
+// probeBalanceAsync logs the site's current free tier after a quota
+// recovery. Best effort — failures only show up at DEBUG.
+func (c *Client) probeBalanceAsync(site config.Site, model config.Model, route Route) {
+	if c.ts == nil {
+		return
+	}
+	gid, ok := c.ts.GuestID(site.Code, route.Proxy)
+	if !ok {
+		return
+	}
+	go func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+		defer cancel()
+		bal, err := c.FetchBalance(ctx, site, gid, model.BotID, route)
+		if err != nil {
+			c.log.Debug("balance probe failed", "site", site.Code, "error", err)
+			return
+		}
+		c.log.Info("site balance", "site", site.Code,
+			"remaining", bal.Free.Remaining, "limit", bal.Free.Limit)
+	}()
+}
+
 // logUpstreamErr records one failed upstream attempt. Quota errors cool the
 // site down and repeats of an identical error inside a short window drop to
 // DEBUG, so a balance-starved site neither taxes every continuation round
@@ -321,12 +358,17 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf
 	routes := c.Routes()
 
 	var lastErr error
+	sawQuota := false
 	for attempt := 0; ; attempt++ {
 		var failed []Route
 		for i, route := range routes {
 			allowSlow := i == 0 && len(routes) > 1
 			err := c.chatOnceSlowStart(ctx, site, modelID, model, prompt, route, allowSlow, info, yield)
 			if err == nil {
+				if sawQuota {
+					c.clearQuotaCool(site.Code)
+					c.probeBalanceAsync(site, model, route)
+				}
 				return nil
 			}
 			lastErr = err
@@ -358,6 +400,9 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf
 			}
 			failed = append(failed, route)
 			c.logUpstreamErr(site.Code, modelID, attempt, route, err)
+			if isQuota(err) {
+				sawQuota = true
+			}
 			if i < len(routes)-1 {
 				continue
 			}
@@ -384,6 +429,15 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf
 			// first cycle retries cheaply and only the second one pays for
 			// a Turnstile re-solve.
 			forceCookie := isSession(lastErr) || (isQuota(lastErr) && attempt >= 1)
+			// The free daily tier is keyed to the visitor cookie: swapping
+			// in a brand-new visitor id restores it instantly, the same
+			// effect as a fresh private window. Only a re-solve (attempt
+			// >= 1) then remains as the slower fallback.
+			if isQuota(lastErr) && attempt == 0 && c.ts != nil {
+				if gid := c.ts.RotateGuest(site.Code); gid != "" {
+					c.log.Info("rotated visitor identity", "site", site.Code, "reason", "quota")
+				}
+			}
 			for _, route := range failed {
 				if err := c.refreshSession(ctx, site, modelID, route, forceCookie); err != nil {
 					c.log.Warn("session refresh failed", "model", modelID, "error", err)

internal/upstream/client_test.go

@@ -10,6 +10,8 @@ import (
 	"time"
 
 	"github.com/nyoungo/dsfree2api/internal/config"
+	"github.com/nyoungo/dsfree2api/internal/httpx"
+	"github.com/nyoungo/dsfree2api/internal/turnstile"
 )
 
 const examplePath = "../../config.example.toml"
@@ -336,3 +338,40 @@ func TestCacheEmptyFailsOverToMirrorWithoutRefreshCycles(t *testing.T) {
 		t.Fatal("mirror was not attempted after cache rejection")
 	}
 }
+
+func TestQuotaRotationRecoversWithoutSolve(t *testing.T) {
+	cfg := testConfig(t)
+	cfg.Upstream.AutoRefresh = true
+	cfg.Upstream.RefreshRetries = 3
+	solver := turnstile.New(cfg.Turnstile, httpx.DefaultUserAgent)
+	const oldGid = "OLDGIDOLDGIDOLDGIDOLDGIDOLDG12"
+	if _, _, err := solver.ImportCookies("de", "cf_clearance=abc; dsgt_gid="+oldGid, ""); err != nil {
+		t.Fatalf("seed cookies: %v", err)
+	}
+	c := New(cfg, solver, nil)
+	calls := 0
+	c.chatOnceOverride = func(_ context.Context, site config.Site, modelID string, model config.Model,
+		_ string, _ Route, _ *ServeInfo, _ func(Event) error) error {
+		calls++
+		if calls == 1 {
+			return newQuota("sse quota exhausted: test")
+		}
+		return nil
+	}
+	if err := c.Chat(context.Background(), "deepseek-v4-flash-de", "p", nil, func(Event) error { return nil }); err != nil {
+		t.Fatalf("Chat: %v", err)
+	}
+	if calls != 2 {
+		t.Fatalf("chat calls = %d, want 2 (quota then success)", calls)
+	}
+	gid, ok := solver.GuestID("de", "")
+	if !ok {
+		t.Fatal("guest id lost after rotation")
+	}
+	if gid == oldGid {
+		t.Fatal("visitor id was not rotated after the quota error")
+	}
+	if c.quotaCooled("de") {
+		t.Error("cooldown should be cleared once the site serves again")
+	}
+}