nyoungo
c877c1bf8c
- 命中额度错误当场调用 RotateGuest（复用既有实现），不再等退避、也不再依赖 auto_refresh；后续路由与下一次请求直接用新访客身份 - quota_notice 信封（type=quota_notice / notice 对象）在通用 error 门之前判定，无 error 字段的额度通知不再被当成静默 delta - 额度文案按站点语言识别：de guthaben/aufgebraucht、es agotado、fr épuisé/consomm、en quota/exhausted，并下钻 data/notice/quota_notice 嵌套字段 - cache message 失败与非 2xx 响应体也按同样规则归类，额度问题会走轮换与跨站 failover 而不是普通错误
1227 lines
33 KiB
Go
1227 lines
33 KiB
Go
package upstream
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"crypto/rand"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"html"
|
||
"io"
|
||
"log/slog"
|
||
"net/http"
|
||
"os"
|
||
"regexp"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
"unicode/utf8"
|
||
|
||
"github.com/nyoungo/dsfree2api/internal/config"
|
||
"github.com/nyoungo/dsfree2api/internal/httpx"
|
||
"github.com/nyoungo/dsfree2api/internal/turnstile"
|
||
)
|
||
|
||
type Route struct {
|
||
Name string
|
||
Proxy string // "" means direct
|
||
}
|
||
|
||
func (r Route) Key() string {
|
||
if r.Proxy == "" {
|
||
return "direct"
|
||
}
|
||
return r.Proxy
|
||
}
|
||
|
||
// RouteCandidate is one egress route supplied by the proxy pool.
|
||
type RouteCandidate struct {
|
||
Name string
|
||
Proxy string
|
||
}
|
||
|
||
// RouteSource resolves per-site route candidates and records request
|
||
// outcomes. Implemented by proxypool.Manager; nil when the pool is unused.
|
||
type RouteSource interface {
|
||
Candidates(site string) []RouteCandidate
|
||
Report(proxy string, ok bool)
|
||
}
|
||
|
||
type chatConfig struct {
|
||
BotID int
|
||
PostID int
|
||
Nonce string
|
||
AJAXURL string
|
||
FetchedAt time.Time
|
||
}
|
||
|
||
type Client struct {
|
||
cfg *config.Config
|
||
ts *turnstile.Solver
|
||
log *slog.Logger
|
||
pool RouteSource
|
||
|
||
stateMu sync.Mutex
|
||
chatCfg map[string]chatConfig
|
||
cfgLocks map[string]*sync.Mutex
|
||
gates map[string]*siteGate
|
||
|
||
// quotaCool marks sites that recently answered with quota errors; they
|
||
// are tried after the healthy mirrors until the cooldown expires.
|
||
quotaCool map[string]time.Time
|
||
// lastWarn* dedupes repeated identical upstream errors so a quota
|
||
// storm cannot drown the log.
|
||
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
|
||
}
|
||
|
||
func New(cfg *config.Config, ts *turnstile.Solver, log *slog.Logger, pool RouteSource) *Client {
|
||
if log == nil {
|
||
log = slog.Default()
|
||
}
|
||
return &Client{
|
||
cfg: cfg,
|
||
ts: ts,
|
||
log: log,
|
||
pool: pool,
|
||
chatCfg: map[string]chatConfig{},
|
||
cfgLocks: map[string]*sync.Mutex{},
|
||
gates: map[string]*siteGate{},
|
||
quotaCool: map[string]time.Time{},
|
||
}
|
||
}
|
||
|
||
// 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()
|
||
primary := strings.TrimSpace(c.cfg.Proxy.URL)
|
||
fallbacks := append([]string(nil), c.cfg.Proxy.FallbackURLs...)
|
||
c.cfg.RUnlock()
|
||
|
||
type cand struct {
|
||
name string
|
||
proxy string
|
||
}
|
||
list := []cand{{"primary", primary}}
|
||
for i, u := range fallbacks {
|
||
list = append(list, cand{fmt.Sprintf("fallback-%d", i+1), strings.TrimSpace(u)})
|
||
}
|
||
seen := map[string]bool{}
|
||
out := make([]Route, 0, len(list))
|
||
for _, it := range list {
|
||
key := it.proxy
|
||
if key == "" {
|
||
key = "direct"
|
||
}
|
||
if seen[key] {
|
||
continue
|
||
}
|
||
seen[key] = true
|
||
out = append(out, Route{Name: it.name, Proxy: it.proxy})
|
||
}
|
||
return out
|
||
}
|
||
|
||
// quotaCoolTTL is how long a site that answered with quota errors keeps
|
||
// being tried after its healthy mirrors.
|
||
const quotaCoolTTL = 10 * time.Minute
|
||
|
||
// quotaCooled reports whether a site is inside its quota cooldown window.
|
||
func (c *Client) quotaCooled(code string) bool {
|
||
c.stateMu.Lock()
|
||
defer c.stateMu.Unlock()
|
||
until, ok := c.quotaCool[code]
|
||
return ok && time.Now().Before(until)
|
||
}
|
||
|
||
// setQuotaCool cools a site down after it reported quota exhaustion.
|
||
// Repeat detections inside the window only extend it; the INFO line marks
|
||
// the transition, not every probe that confirms it.
|
||
func (c *Client) setQuotaCool(code string) {
|
||
c.stateMu.Lock()
|
||
_, wasCooled := c.quotaCool[code]
|
||
// 每次都把截止时间顺延，重复探测只延长窗口；只有从"未冷却"进入
|
||
// "冷却"这一次状态翻转才记 INFO，之后的确认探测不再刷日志。
|
||
c.quotaCool[code] = time.Now().Add(quotaCoolTTL)
|
||
c.stateMu.Unlock()
|
||
if !wasCooled {
|
||
c.log.Info("site quota cooldown", "site", code, "for", quotaCoolTTL)
|
||
}
|
||
}
|
||
|
||
// sortQuotaCoolLast moves quota-cooled sites behind the healthy candidates,
|
||
// preserving relative order inside both groups.
|
||
func (c *Client) sortQuotaCoolLast(ids []string) []string {
|
||
var cooled, rest []string
|
||
for _, id := range ids {
|
||
m, ok := c.modelCopy(id)
|
||
if ok && c.quotaCooled(m.Site) {
|
||
cooled = append(cooled, id)
|
||
continue
|
||
}
|
||
rest = append(rest, id)
|
||
}
|
||
if len(cooled) == 0 {
|
||
return ids
|
||
}
|
||
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
|
||
// nor floods the log.
|
||
func (c *Client) logUpstreamErr(site, modelID string, attempt int, route Route, err error) {
|
||
if isQuota(err) {
|
||
c.setQuotaCool(site)
|
||
}
|
||
msg := err.Error()
|
||
c.stateMu.Lock()
|
||
dup := msg == c.lastWarnMsg && time.Since(c.lastWarnAt) < 30*time.Second
|
||
c.lastWarnMsg = msg
|
||
c.lastWarnAt = time.Now()
|
||
c.stateMu.Unlock()
|
||
if dup {
|
||
c.log.Debug("upstream error (repeating)", "model", modelID, "attempt", attempt+1, "route", route.Name, "error", err)
|
||
return
|
||
}
|
||
c.log.Warn("upstream error", "model", modelID, "attempt", attempt+1, "route", route.Name, "error", err)
|
||
}
|
||
|
||
// logCacheReject fingerprints a prompt the site's cache step called empty:
|
||
// byte class counts plus head/tail quotes tell whether the rejection tracks
|
||
// prompt content (WAF/filter) or session state.
|
||
func (c *Client) logCacheReject(site, modelID, prompt string) {
|
||
nul, ctrl := 0, 0
|
||
for i := 0; i < len(prompt); i++ {
|
||
b := prompt[i]
|
||
switch {
|
||
case b == 0:
|
||
nul++
|
||
case b < 0x20 && b != '\n' && b != '\r' && b != '\t':
|
||
ctrl++
|
||
}
|
||
}
|
||
head, tail := prompt, ""
|
||
if r := []rune(prompt); len(r) > 160 {
|
||
head, tail = string(r[:80]), string(r[len(r)-80:])
|
||
}
|
||
sum := sha256.Sum256([]byte(prompt))
|
||
c.log.Warn("cache message rejected",
|
||
"site", site, "model", modelID, "chars", len(prompt),
|
||
"nul", nul, "ctrl", ctrl, "sha", hex.EncodeToString(sum[:8]),
|
||
"head", strconv.Quote(head), "tail", strconv.Quote(tail))
|
||
}
|
||
|
||
// routesFor returns the ordered candidate routes for one site: proxy-pool
|
||
// candidates first (sticky, health filtered), then the global primary and
|
||
// fallback lines as the safety net.
|
||
func (c *Client) routesFor(siteCode string) []Route {
|
||
out := make([]Route, 0, 4)
|
||
seen := map[string]bool{}
|
||
add := func(r Route) {
|
||
key := r.Key()
|
||
if seen[key] {
|
||
return
|
||
}
|
||
seen[key] = true
|
||
out = append(out, r)
|
||
}
|
||
if c.pool != nil {
|
||
for _, cand := range c.pool.Candidates(siteCode) {
|
||
add(Route{Name: cand.Name, Proxy: cand.Proxy})
|
||
}
|
||
}
|
||
for _, r := range c.Routes() {
|
||
add(r)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// poolReport feeds a route outcome back to the proxy pool; it is a no-op for
|
||
// global routes and the direct connection.
|
||
func (c *Client) poolReport(proxy string, ok bool) {
|
||
if c.pool == nil || proxy == "" {
|
||
return
|
||
}
|
||
c.pool.Report(proxy, ok)
|
||
}
|
||
|
||
// Chat streams the prompt through the upstream site, honouring route
|
||
// failover, slow-start detection, session refresh and cross-site failover.
|
||
func (c *Client) Chat(ctx context.Context, modelID string, prompt string, info *ServeInfo, yield func(Event) error) error {
|
||
model, ok := c.modelCopy(modelID)
|
||
if !ok {
|
||
return errf("unknown model %s", modelID)
|
||
}
|
||
if !model.Enabled {
|
||
return errf("model %s is disabled", modelID)
|
||
}
|
||
|
||
candidates := []string{modelID}
|
||
if c.cfg.Upstream.CrossSiteFailover {
|
||
for _, alt := range c.alternates(modelID, model) {
|
||
candidates = append(candidates, alt)
|
||
}
|
||
}
|
||
// A site that recently reported quota errors keeps being tried, but
|
||
// behind the healthy mirrors — continuation rounds stop paying its
|
||
// refresh tax while its balance is empty.
|
||
candidates = c.sortQuotaCoolLast(candidates)
|
||
|
||
var last error
|
||
triedPrimary := false
|
||
prev := ""
|
||
for _, id := range candidates {
|
||
m := model
|
||
if id == modelID {
|
||
triedPrimary = true
|
||
} else {
|
||
mc, ok := c.modelCopy(id)
|
||
if !ok || !mc.Enabled {
|
||
continue
|
||
}
|
||
m = mc
|
||
}
|
||
// Only a real switch after a failed candidate is a failover; the
|
||
// first pick of a request (even a sorted-past primary) is not.
|
||
if prev != "" && prev != id {
|
||
c.log.Warn("cross-site failover", "from", prev, "to", id)
|
||
}
|
||
err := c.chatWithRetries(ctx, id, m, prompt, info, yield)
|
||
if err == nil {
|
||
return nil
|
||
}
|
||
last = err
|
||
if contentStarted(err) || !turnstile.Retryable(err) {
|
||
return AsUpstream(err)
|
||
}
|
||
// Quota means this site has no balance and a cache-rejection means
|
||
// its session is off: the next candidate is the point. Other errors
|
||
// keep going only while the primary has not had its chance yet (a
|
||
// quota-cooled primary sits at the end of the list); once it failed
|
||
// too, report that error as before.
|
||
if !isQuota(err) && !isCacheEmptyProblem(err) && triedPrimary {
|
||
return AsUpstream(err)
|
||
}
|
||
prev = id
|
||
}
|
||
return AsUpstream(last)
|
||
}
|
||
|
||
// alternates lists other enabled sites hosting the same upstream model.
|
||
func (c *Client) alternates(modelID string, m config.Model) []string {
|
||
c.cfg.RLock()
|
||
defer c.cfg.RUnlock()
|
||
var out []string
|
||
for id, other := range c.cfg.Models {
|
||
if id == modelID || other.Site == m.Site {
|
||
continue
|
||
}
|
||
if other.UpstreamID != m.UpstreamID || !other.Enabled {
|
||
continue
|
||
}
|
||
if s, ok := c.cfg.Sites[other.Site]; !ok || !s.Enabled {
|
||
continue
|
||
}
|
||
out = append(out, id)
|
||
}
|
||
return out
|
||
}
|
||
|
||
func (c *Client) modelCopy(id string) (config.Model, bool) {
|
||
c.cfg.RLock()
|
||
defer c.cfg.RUnlock()
|
||
m, ok := c.cfg.Models[id]
|
||
if !ok {
|
||
return config.Model{}, false
|
||
}
|
||
return *m, true
|
||
}
|
||
|
||
func (c *Client) siteCopy(code string) (config.Site, bool) {
|
||
c.cfg.RLock()
|
||
defer c.cfg.RUnlock()
|
||
s, ok := c.cfg.Sites[code]
|
||
if !ok {
|
||
return config.Site{}, false
|
||
}
|
||
return *s, true
|
||
}
|
||
|
||
func (c *Client) userAgent() string {
|
||
c.cfg.RLock()
|
||
defer c.cfg.RUnlock()
|
||
return c.cfg.Upstream.UserAgent
|
||
}
|
||
|
||
func (c *Client) chatWithRetries(ctx context.Context, modelID string, model config.Model, prompt string, info *ServeInfo, yield func(Event) error) error {
|
||
site, ok := c.siteCopy(model.Site)
|
||
if !ok {
|
||
return errf("unknown site %s", model.Site)
|
||
}
|
||
if !site.Enabled {
|
||
return errf("site %s is disabled", model.Site)
|
||
}
|
||
routes := c.routesFor(site.Code)
|
||
quotaNotified := false
|
||
|
||
var lastErr error
|
||
sawQuota := false
|
||
rotatedQuota := false
|
||
for attempt := 0; ; attempt++ {
|
||
var failed []Route
|
||
for i, route := range routes {
|
||
// A first-time Turnstile solve can far exceed the slow-start
|
||
// budget; only arm the timer when cookies are already cached
|
||
// (or no solver is configured at all).
|
||
cookiesReady := c.ts == nil || c.ts.CookiesReady(site.Code, route.Proxy)
|
||
allowSlow := i == 0 && len(routes) > 1 && cookiesReady
|
||
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)
|
||
}
|
||
c.poolReport(route.Proxy, true)
|
||
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)
|
||
}
|
||
// A solver-side rejection (bad key, rejected challenge) will fail
|
||
// identically on every attempt — stop instead of burning retries.
|
||
if !turnstile.Retryable(err) {
|
||
return AsUpstream(err)
|
||
}
|
||
// The site demands Turnstile verification while the solver is
|
||
// switched off: refreshing cookies cannot help, say why instead.
|
||
if isSession(err) && (c.ts == nil || !c.ts.Enabled()) {
|
||
return &UpstreamError{
|
||
Msg: "site requires Turnstile verification — import cookies on the console " +
|
||
"Turnstile page, or enable [turnstile] to solve automatically",
|
||
Cause: err,
|
||
}
|
||
}
|
||
if isConfigProblem(err) {
|
||
return AsUpstream(err)
|
||
}
|
||
// The cache step rejected the message in a way config drops and
|
||
// cookie refreshes do not fix; stop burning refresh cycles so
|
||
// the caller can fail over to a mirror right away.
|
||
if isCacheEmptyProblem(err) {
|
||
return AsUpstream(err)
|
||
}
|
||
// Quota/session failures are site-side, not proxy-side — do not
|
||
// blame the egress endpoint for them.
|
||
if !isQuota(err) && !isSession(err) {
|
||
c.poolReport(route.Proxy, false)
|
||
}
|
||
failed = append(failed, route)
|
||
c.logUpstreamErr(site.Code, modelID, attempt, route, err)
|
||
if isQuota(err) {
|
||
sawQuota = true
|
||
// The free daily tier is keyed to the visitor cookie, so a
|
||
// brand-new visitor id restores it instantly — the same
|
||
// effect as a fresh private window. Rotate the moment the
|
||
// site reports exhaustion, before the backoff wait and
|
||
// before the next route/attempt, so the retry (and any
|
||
// sibling route tried right after) is billed to a balance
|
||
// that is not already spent. This also runs when
|
||
// auto_refresh is off, leaving the new identity in place
|
||
// for the next request instead of the expired one.
|
||
if !rotatedQuota && c.ts != nil {
|
||
rotatedQuota = true
|
||
if gid := c.ts.RotateGuest(site.Code); gid != "" {
|
||
c.log.Info("rotated visitor identity", "site", site.Code, "reason", "quota")
|
||
}
|
||
}
|
||
}
|
||
if i < len(routes)-1 {
|
||
continue
|
||
}
|
||
break
|
||
}
|
||
|
||
c.cfg.RLock()
|
||
autoRefresh := c.cfg.Upstream.AutoRefresh
|
||
backoff := c.cfg.Upstream.RetryBackoffSeconds
|
||
retries := c.cfg.Upstream.RefreshRetries
|
||
c.cfg.RUnlock()
|
||
if !autoRefresh || attempt >= retries {
|
||
break
|
||
}
|
||
if len(failed) > 0 {
|
||
wait := time.Duration(float64(attempt+1) * backoff * float64(time.Second))
|
||
select {
|
||
case <-ctx.Done():
|
||
return AsUpstream(ctx.Err())
|
||
case <-time.After(wait):
|
||
}
|
||
// Session errors need a fresh cookie right away. Quota errors
|
||
// often clear once the cached chat config is dropped, so the
|
||
// first cycle retries cheaply and only the second one pays for
|
||
// a Turnstile re-solve. The visitor id was already swapped the
|
||
// moment exhaustion was detected; a full re-solve (attempt >= 1)
|
||
// is the slower fallback for when that alone was not enough.
|
||
forceCookie := isSession(lastErr) || (isQuota(lastErr) && attempt >= 1)
|
||
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)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return AsUpstream(lastErr)
|
||
}
|
||
|
||
// chatOnceSlowStart bounds the time the primary route may take to produce its
|
||
// first upstream event before a fallback route is tried instead. Errors that
|
||
// happen after the first event are marked ContentStarted so callers never
|
||
// restart a response the client already partially received.
|
||
func (c *Client) chatOnceSlowStart(
|
||
ctx context.Context,
|
||
site config.Site,
|
||
modelID string,
|
||
model config.Model,
|
||
prompt string,
|
||
route Route,
|
||
allowSlow bool,
|
||
info *ServeInfo,
|
||
yield func(Event) error,
|
||
) error {
|
||
ctx, cancel := context.WithCancel(ctx)
|
||
defer cancel()
|
||
|
||
events := make(chan Event, 16)
|
||
errCh := make(chan error, 1)
|
||
gotFirst := false
|
||
|
||
go func() {
|
||
fn := c.chatOnce
|
||
if c.chatOnceOverride != nil {
|
||
fn = c.chatOnceOverride
|
||
}
|
||
err := fn(ctx, site, modelID, model, prompt, route, info, func(ev Event) error {
|
||
select {
|
||
case events <- ev:
|
||
return nil
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
}
|
||
})
|
||
close(events)
|
||
errCh <- err
|
||
}()
|
||
|
||
mark := func(err error) error {
|
||
if err == nil || !gotFirst {
|
||
return err
|
||
}
|
||
ue := AsUpstream(err)
|
||
ue.ContentStarted = true
|
||
return ue
|
||
}
|
||
take := func(ev Event) error {
|
||
gotFirst = true
|
||
if err := yield(ev); err != nil {
|
||
return mark(err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
if allowSlow {
|
||
var slow float64
|
||
c.cfg.RLock()
|
||
slow = c.cfg.Proxy.SlowStartSeconds
|
||
c.cfg.RUnlock()
|
||
if slow > 0 {
|
||
timer := time.NewTimer(time.Duration(slow * float64(time.Second)))
|
||
defer timer.Stop()
|
||
select {
|
||
case ev, ok := <-events:
|
||
if !ok {
|
||
return mark(<-errCh)
|
||
}
|
||
if err := take(ev); err != nil {
|
||
return err
|
||
}
|
||
case err := <-errCh:
|
||
for ev := range events {
|
||
if e := take(ev); e != nil {
|
||
return e
|
||
}
|
||
}
|
||
return mark(err)
|
||
case <-timer.C:
|
||
cancel()
|
||
<-errCh
|
||
return errf("primary route produced no upstream event within %gs", slow)
|
||
}
|
||
}
|
||
}
|
||
|
||
for {
|
||
ev, ok := <-events
|
||
if !ok {
|
||
return mark(<-errCh)
|
||
}
|
||
if err := take(ev); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Client) chatOnce(
|
||
ctx context.Context,
|
||
site config.Site,
|
||
modelID string,
|
||
model config.Model,
|
||
prompt string,
|
||
route Route,
|
||
info *ServeInfo,
|
||
yield func(Event) error,
|
||
) error {
|
||
if err := c.acquire(ctx, site.Code); err != nil {
|
||
return AsUpstream(err)
|
||
}
|
||
defer c.release(site.Code)
|
||
|
||
if info != nil {
|
||
info.Set(modelID, site.Code, route.Name)
|
||
}
|
||
|
||
c.cfg.RLock()
|
||
streamTimeout := time.Duration(c.cfg.Upstream.StreamTimeout * float64(time.Second))
|
||
timeout := time.Duration(c.cfg.Upstream.Timeout * float64(time.Second))
|
||
ua := c.cfg.Upstream.UserAgent
|
||
c.cfg.RUnlock()
|
||
|
||
sess, err := httpx.NewSession(httpx.Options{ProxyURL: route.Proxy, Timeout: streamTimeout})
|
||
if err != nil {
|
||
return errf("create session: %v", err)
|
||
}
|
||
defer sess.Close()
|
||
|
||
if err := c.ts.ApplyValidCookies(ctx, sess, &site, route.Proxy); err != nil {
|
||
return AsUpstream(err)
|
||
}
|
||
|
||
cc, err := c.chatConfig(ctx, sess, site, modelID, model, route, ua)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
referer := site.BaseURL + model.PagePath
|
||
sessionID := newUUID()
|
||
convUUID := newUUID()
|
||
clientMsgID := fmt.Sprintf("aipkit-client-msg-%d-%d-%s", cc.BotID, time.Now().UnixMilli(), randHex(3))
|
||
|
||
resp, err := sess.Do(ctx, httpx.Request{
|
||
Method: "POST",
|
||
URL: cc.AJAXURL,
|
||
Header: ajaxHeaders(site, referer, ua),
|
||
Body: []byte(encodeForm(map[string]string{
|
||
"action": "aipkit_cache_sse_message",
|
||
"message": prompt,
|
||
"_ajax_nonce": cc.Nonce,
|
||
"bot_id": itoa(cc.BotID),
|
||
"session_id": sessionID,
|
||
"conversation_uuid": convUUID,
|
||
"user_client_message_id": clientMsgID,
|
||
})),
|
||
Timeout: timeout,
|
||
})
|
||
if err != nil {
|
||
return errf("cache message request failed: %v", err)
|
||
}
|
||
if bytes.Contains(resp.Body, []byte("empty_data_to_cache")) {
|
||
c.logCacheReject(site.Code, modelID, prompt)
|
||
}
|
||
if resp.Status >= 400 {
|
||
if resp.Status == 403 || resp.Status == 401 {
|
||
return newSession(fmt.Sprintf("cache message http %d", resp.Status))
|
||
}
|
||
// The allowance check can trip at the cache step too (the site
|
||
// counts the prompt against the daily total before streaming), and
|
||
// the mirrors word it per locale — classify it instead of letting
|
||
// it surface as a generic error that neither rotates nor fails over.
|
||
if p := jsonPayload(resp.Body); isQuotaPayload(p) {
|
||
return newQuota("cache message quota exhausted: " + truncate(string(resp.Body), 400))
|
||
}
|
||
return errf("cache message http %d: %s", resp.Status, truncate(string(resp.Body), 400))
|
||
}
|
||
|
||
var init struct {
|
||
Success bool `json:"success"`
|
||
Data json.RawMessage `json:"data"`
|
||
}
|
||
if err := json.Unmarshal(resp.Body, &init); err != nil {
|
||
return errf("cache message bad json: %v", err)
|
||
}
|
||
if !init.Success {
|
||
if code := extractCode(init.Data); code != "" && strings.Contains(strings.ToLower(code), "nonce") {
|
||
c.dropChatCfg(modelID, route)
|
||
}
|
||
payload := jsonPayload(resp.Body)
|
||
if isQuotaPayload(payload) {
|
||
return newQuota("cache message quota exhausted: " + truncate(string(resp.Body), 400))
|
||
}
|
||
if isSessionPayload(payload) {
|
||
return newSession("cache message session failed: " + truncate(string(resp.Body), 400))
|
||
}
|
||
return errf("cache message failed: %s", truncate(string(resp.Body), 400))
|
||
}
|
||
var dataObj struct {
|
||
CacheKey string `json:"cache_key"`
|
||
}
|
||
if err := json.Unmarshal(init.Data, &dataObj); err != nil || dataObj.CacheKey == "" {
|
||
return errf("cache message missing cache_key: %s", truncate(string(resp.Body), 400))
|
||
}
|
||
|
||
params := encodeForm(map[string]string{
|
||
"action": "aipkit_frontend_chat_stream",
|
||
"cache_key": dataObj.CacheKey,
|
||
"bot_id": itoa(cc.BotID),
|
||
"session_id": sessionID,
|
||
"conversation_uuid": convUUID,
|
||
"post_id": itoa(cc.PostID),
|
||
"_ts": itoa64(time.Now().UnixMilli()),
|
||
"_ajax_nonce": cc.Nonce,
|
||
})
|
||
|
||
streamCtx, streamCancel := context.WithTimeout(ctx, streamTimeout)
|
||
defer streamCancel()
|
||
|
||
headers := ajaxHeaders(site, referer, ua)
|
||
headers["Accept"] = "text/event-stream"
|
||
|
||
rc, status, err := sess.Stream(streamCtx, httpx.Request{
|
||
Method: "GET",
|
||
URL: cc.AJAXURL + "?" + params,
|
||
Header: headers,
|
||
})
|
||
if err != nil {
|
||
if streamCtx.Err() != nil && ctx.Err() == nil {
|
||
return errf("upstream stream timeout: %v", err)
|
||
}
|
||
return streamHTTPError(status, err)
|
||
}
|
||
defer rc.Close()
|
||
|
||
if os.Getenv("DSFREE_SSE_TIMING") != "" {
|
||
start := time.Now()
|
||
rc = &sseTimingReader{rc: rc, start: start, logf: func(n int) {
|
||
c.log.Info("sse read", "at_ms", time.Since(start).Milliseconds(), "n", n)
|
||
}}
|
||
}
|
||
|
||
return readSSE(rc, func(event string, lines []string) error {
|
||
events, err := translateEvent(event, lines)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
for _, ev := range events {
|
||
if err := yield(ev); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
})
|
||
}
|
||
|
||
// streamHTTPError classifies a stream setup failure by its HTTP status. 401/403
|
||
// mean the cached session is no longer accepted (expired nonce / Turnstile
|
||
// cookie), which must be reported as a session error so the caller refreshes
|
||
// cookies instead of blaming the egress endpoint (poolReport(false)) and
|
||
// keeping the stale cookie.
|
||
func streamHTTPError(status int, err error) error {
|
||
if status == http.StatusUnauthorized || status == http.StatusForbidden {
|
||
return newSession(fmt.Sprintf("upstream stream http %d: %v", status, err))
|
||
}
|
||
return AsUpstream(err)
|
||
}
|
||
|
||
// sseTimingReader (diagnosis only, gated by DSFREE_SSE_TIMING) logs when bytes
|
||
// arrive from the upstream stream.
|
||
type sseTimingReader struct {
|
||
rc io.ReadCloser
|
||
start time.Time
|
||
logf func(n int)
|
||
}
|
||
|
||
func (r *sseTimingReader) Read(p []byte) (int, error) {
|
||
n, err := r.rc.Read(p)
|
||
if n > 0 {
|
||
r.logf(n)
|
||
}
|
||
return n, err
|
||
}
|
||
|
||
func (r *sseTimingReader) Close() error { return r.rc.Close() }
|
||
|
||
// chatConfig fetches (and caches) the page-embedded nonce for a model.
|
||
func (c *Client) chatConfig(
|
||
ctx context.Context,
|
||
sess httpx.Session,
|
||
site config.Site,
|
||
modelID string,
|
||
model config.Model,
|
||
route Route,
|
||
ua string,
|
||
) (chatConfig, error) {
|
||
key := modelID + "|" + route.Key()
|
||
if v, ok := c.loadChatCfg(key); ok {
|
||
return v, nil
|
||
}
|
||
lock := c.lockFor(key)
|
||
lock.Lock()
|
||
defer lock.Unlock()
|
||
if v, ok := c.loadChatCfg(key); ok {
|
||
return v, nil
|
||
}
|
||
|
||
pageURL := site.BaseURL + model.PagePath
|
||
resp, err := sess.Do(ctx, httpx.Request{
|
||
Method: "GET",
|
||
URL: pageURL,
|
||
Header: browserHeaders(site, ua),
|
||
Timeout: 60 * time.Second,
|
||
})
|
||
if err != nil {
|
||
return chatConfig{}, errf("chat config fetch failed: %v", err)
|
||
}
|
||
if resp.Status >= 400 {
|
||
return chatConfig{}, errf("chat config http %d", resp.Status)
|
||
}
|
||
|
||
parsed, found := parseChatConfig(string(resp.Body), site, model)
|
||
if !found {
|
||
if model.BotID <= 0 || model.PostID <= 0 {
|
||
return chatConfig{}, errf("no chat data-config found for %s", modelID)
|
||
}
|
||
nonce, err := c.fetchNonce(ctx, sess, site, model, ua)
|
||
if err != nil {
|
||
return chatConfig{}, err
|
||
}
|
||
parsed = chatConfig{
|
||
BotID: model.BotID, PostID: model.PostID, Nonce: nonce,
|
||
AJAXURL: site.AJAXURL, FetchedAt: time.Now(),
|
||
}
|
||
}
|
||
parsed.FetchedAt = time.Now()
|
||
c.storeChatCfg(key, parsed)
|
||
c.log.Info("loaded chat config", "model", modelID, "route", route.Name,
|
||
"bot_id", parsed.BotID, "post_id", parsed.PostID)
|
||
return parsed, nil
|
||
}
|
||
|
||
var containerRe = regexp.MustCompile(`id="aipkit_chat_container_(\d+)"[^>]*data-config='([^']+)'`)
|
||
|
||
func parseChatConfig(page string, site config.Site, model config.Model) (chatConfig, bool) {
|
||
for _, m := range containerRe.FindAllStringSubmatch(page, -1) {
|
||
botID := atoi(m[1])
|
||
raw := html.UnescapeString(m[2])
|
||
var cfg struct {
|
||
BotID any `json:"botId"`
|
||
PostID any `json:"postId"`
|
||
Nonce string `json:"nonce"`
|
||
AjaxURL string `json:"ajaxUrl"`
|
||
}
|
||
if err := json.Unmarshal([]byte(raw), &cfg); err != nil {
|
||
continue
|
||
}
|
||
if model.BotID > 0 && botID != model.BotID {
|
||
continue
|
||
}
|
||
out := chatConfig{
|
||
BotID: toInt(cfg.BotID, botID),
|
||
PostID: toInt(cfg.PostID, model.PostID),
|
||
Nonce: cfg.Nonce,
|
||
AJAXURL: cfg.AjaxURL,
|
||
FetchedAt: time.Now(),
|
||
}
|
||
if out.AJAXURL == "" {
|
||
out.AJAXURL = site.AJAXURL
|
||
}
|
||
if out.Nonce != "" {
|
||
return out, true
|
||
}
|
||
}
|
||
return chatConfig{}, false
|
||
}
|
||
|
||
func (c *Client) fetchNonce(ctx context.Context, sess httpx.Session, site config.Site, model config.Model, ua string) (string, error) {
|
||
resp, err := sess.Do(ctx, httpx.Request{
|
||
Method: "POST",
|
||
URL: site.AJAXURL,
|
||
Header: ajaxHeaders(site, site.BaseURL+model.PagePath, ua),
|
||
Body: []byte(encodeForm(map[string]string{"action": "aipkit_get_frontend_chat_nonce", "bot_id": itoa(model.BotID)})),
|
||
Timeout: 30 * time.Second,
|
||
})
|
||
if err != nil {
|
||
return "", errf("nonce fetch failed: %v", err)
|
||
}
|
||
var out struct {
|
||
Success bool `json:"success"`
|
||
Data struct {
|
||
Nonce string `json:"nonce"`
|
||
} `json:"data"`
|
||
}
|
||
if err := json.Unmarshal(resp.Body, &out); err != nil || !out.Success || out.Data.Nonce == "" {
|
||
return "", errf("nonce fetch failed for %s: %s", model.ID, truncate(string(resp.Body), 400))
|
||
}
|
||
return out.Data.Nonce, nil
|
||
}
|
||
|
||
// refreshSession drops cached state for a route and optionally re-solves Turnstile.
|
||
func (c *Client) refreshSession(
|
||
ctx context.Context,
|
||
site config.Site,
|
||
modelID string,
|
||
route Route,
|
||
invalidateCookies bool,
|
||
) error {
|
||
c.dropChatCfg(modelID, route)
|
||
if !invalidateCookies {
|
||
return nil
|
||
}
|
||
c.log.Info("invalidate verified cookie cache", "site", site.Code, "route", route.Name)
|
||
c.ts.Invalidate(site.Code)
|
||
|
||
sess, err := httpx.NewSession(httpx.Options{ProxyURL: route.Proxy, Timeout: 60 * time.Second})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer sess.Close()
|
||
return c.ts.Refresh(ctx, sess, &site, route.Proxy)
|
||
}
|
||
|
||
// InvalidateAll clears cached nonces and cookies (web console action).
|
||
func (c *Client) InvalidateAll() {
|
||
c.stateMu.Lock()
|
||
c.chatCfg = map[string]chatConfig{}
|
||
c.stateMu.Unlock()
|
||
c.cfg.RLock()
|
||
codes := make([]string, 0, len(c.cfg.Sites))
|
||
for code := range c.cfg.Sites {
|
||
codes = append(codes, code)
|
||
}
|
||
c.cfg.RUnlock()
|
||
for _, code := range codes {
|
||
c.ts.Invalidate(code)
|
||
}
|
||
}
|
||
|
||
func (c *Client) ConfigCacheSize() int {
|
||
c.stateMu.Lock()
|
||
defer c.stateMu.Unlock()
|
||
return len(c.chatCfg)
|
||
}
|
||
|
||
// ── cache helpers ────────────────────────────────────────────────
|
||
|
||
func (c *Client) loadChatCfg(key string) (chatConfig, bool) {
|
||
c.stateMu.Lock()
|
||
defer c.stateMu.Unlock()
|
||
v, ok := c.chatCfg[key]
|
||
if !ok {
|
||
return chatConfig{}, false
|
||
}
|
||
c.cfg.RLock()
|
||
ttl := time.Duration(c.cfg.Upstream.ConfigTTLSeconds) * time.Second
|
||
c.cfg.RUnlock()
|
||
if ttl <= 0 || time.Since(v.FetchedAt) > ttl {
|
||
return chatConfig{}, false
|
||
}
|
||
return v, true
|
||
}
|
||
|
||
func (c *Client) storeChatCfg(key string, v chatConfig) {
|
||
c.stateMu.Lock()
|
||
defer c.stateMu.Unlock()
|
||
c.chatCfg[key] = v
|
||
}
|
||
|
||
func (c *Client) dropChatCfg(modelID string, route Route) {
|
||
suffix := "|" + route.Key()
|
||
c.stateMu.Lock()
|
||
defer c.stateMu.Unlock()
|
||
for k := range c.chatCfg {
|
||
if strings.HasSuffix(k, suffix) && strings.TrimSuffix(k, suffix) == modelID {
|
||
delete(c.chatCfg, k)
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Client) lockFor(key string) *sync.Mutex {
|
||
c.stateMu.Lock()
|
||
defer c.stateMu.Unlock()
|
||
l, ok := c.cfgLocks[key]
|
||
if !ok {
|
||
l = &sync.Mutex{}
|
||
c.cfgLocks[key] = l
|
||
}
|
||
return l
|
||
}
|
||
|
||
// ── concurrency gate ────────────────────────────────────────────
|
||
|
||
// siteGate 是单个站点的并发闸门。名额用 busy 计数而不是"首次使用时定容的
|
||
// channel"，这样 [limits].max_concurrent_per_site 被修改后立即生效，无需重启
|
||
// （channel 版本只在第一次 acquire 时定容，之后改配置一直不生效）。
|
||
type siteGate struct {
|
||
mu sync.Mutex
|
||
busy int
|
||
// wake 在每次 release 时被关闭并重建，用来广播唤醒等待者；等待者醒来后
|
||
// 按当时最新的配置重新竞争名额。
|
||
wake chan struct{}
|
||
}
|
||
|
||
func newSiteGate() *siteGate { return &siteGate{wake: make(chan struct{})} }
|
||
|
||
func (c *Client) acquire(ctx context.Context, site string) error {
|
||
c.stateMu.Lock()
|
||
g, ok := c.gates[site]
|
||
if !ok {
|
||
g = newSiteGate()
|
||
c.gates[site] = g
|
||
}
|
||
c.stateMu.Unlock()
|
||
return c.admit(ctx, site, g)
|
||
}
|
||
|
||
// admit 在当前配置的上限内占一个名额；超限时阻塞到有空位或 ctx 结束。
|
||
func (c *Client) admit(ctx context.Context, site string, g *siteGate) error {
|
||
for {
|
||
c.cfg.RLock()
|
||
max := c.cfg.Limits.MaxConcurrentPerSite
|
||
c.cfg.RUnlock()
|
||
|
||
g.mu.Lock()
|
||
if max <= 0 || g.busy < max {
|
||
g.busy++
|
||
g.mu.Unlock()
|
||
return nil
|
||
}
|
||
wake := g.wake
|
||
g.mu.Unlock()
|
||
|
||
select {
|
||
case <-ctx.Done():
|
||
return fmt.Errorf("site %s busy: %w", site, ctx.Err())
|
||
case <-wake:
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Client) release(site string) {
|
||
c.stateMu.Lock()
|
||
g, ok := c.gates[site]
|
||
c.stateMu.Unlock()
|
||
if !ok {
|
||
return
|
||
}
|
||
g.mu.Lock()
|
||
if g.busy > 0 {
|
||
g.busy--
|
||
}
|
||
close(g.wake)
|
||
g.wake = make(chan struct{})
|
||
g.mu.Unlock()
|
||
}
|
||
|
||
// ── helpers ─────────────────────────────────────────────────────
|
||
|
||
func browserHeaders(site config.Site, ua string) map[string]string {
|
||
if ua == "" {
|
||
ua = httpx.DefaultUserAgent
|
||
}
|
||
return map[string]string{
|
||
"User-Agent": ua,
|
||
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
|
||
"Accept-Language": site.Language,
|
||
"Origin": site.BaseURL,
|
||
"Referer": site.BaseURL + "/",
|
||
}
|
||
}
|
||
|
||
func ajaxHeaders(site config.Site, referer, ua string) map[string]string {
|
||
if ua == "" {
|
||
ua = httpx.DefaultUserAgent
|
||
}
|
||
return map[string]string{
|
||
"User-Agent": ua,
|
||
"Accept": "application/json, text/javascript, */*; q=0.01",
|
||
"Accept-Language": site.Language,
|
||
"X-Requested-With": "XMLHttpRequest",
|
||
"Origin": site.BaseURL,
|
||
"Referer": referer,
|
||
}
|
||
}
|
||
|
||
func encodeForm(v map[string]string) string {
|
||
keys := make([]string, 0, len(v))
|
||
for k := range v {
|
||
keys = append(keys, k)
|
||
}
|
||
sortStrings(keys)
|
||
var b strings.Builder
|
||
for i, k := range keys {
|
||
if i > 0 {
|
||
b.WriteByte('&')
|
||
}
|
||
b.WriteString(urlEncode(k))
|
||
b.WriteByte('=')
|
||
b.WriteString(urlEncode(v[k]))
|
||
}
|
||
return b.String()
|
||
}
|
||
|
||
func urlEncode(s string) string {
|
||
var b strings.Builder
|
||
for i := 0; i < len(s); i++ {
|
||
ch := s[i]
|
||
switch {
|
||
case ch >= 'a' && ch <= 'z', ch >= 'A' && ch <= 'Z', ch >= '0' && ch <= '9',
|
||
ch == '-', ch == '_', ch == '.', ch == '~':
|
||
b.WriteByte(ch)
|
||
case ch == ' ':
|
||
b.WriteByte('+')
|
||
default:
|
||
fmt.Fprintf(&b, "%%%02X", ch)
|
||
}
|
||
}
|
||
return b.String()
|
||
}
|
||
|
||
func sortStrings(s []string) {
|
||
for i := 1; i < len(s); i++ {
|
||
for j := i; j > 0 && s[j] < s[j-1]; j-- {
|
||
s[j], s[j-1] = s[j-1], s[j]
|
||
}
|
||
}
|
||
}
|
||
|
||
func newUUID() string { return randHex(16) }
|
||
|
||
func randHex(n int) string {
|
||
b := make([]byte, n)
|
||
if _, err := rand.Read(b); err != nil {
|
||
return strings.Repeat("0", n*2)
|
||
}
|
||
return hex.EncodeToString(b)
|
||
}
|
||
|
||
func itoa(v int) string { return fmt.Sprintf("%d", v) }
|
||
func itoa64(v int64) string { return fmt.Sprintf("%d", v) }
|
||
|
||
func truncate(s string, n int) string {
|
||
if len(s) <= n {
|
||
return s
|
||
}
|
||
// 按字节截断可能切碎多字节字符，回退到字符边界，避免产出非法 UTF-8。
|
||
for n > 0 && !utf8.RuneStart(s[n]) {
|
||
n--
|
||
}
|
||
return s[:n] + "..."
|
||
}
|
||
|
||
func atoi(s string) int {
|
||
n := 0
|
||
for _, ch := range s {
|
||
if ch < '0' || ch > '9' {
|
||
return n
|
||
}
|
||
n = n*10 + int(ch-'0')
|
||
}
|
||
return n
|
||
}
|
||
|
||
func extractCode(data json.RawMessage) string {
|
||
var obj struct {
|
||
Code string `json:"code"`
|
||
}
|
||
if err := json.Unmarshal(data, &obj); err == nil {
|
||
return obj.Code
|
||
}
|
||
return ""
|
||
}
|
||
|
||
func toInt(v any, def int) int {
|
||
switch t := v.(type) {
|
||
case float64:
|
||
return int(t)
|
||
case int:
|
||
return t
|
||
case string:
|
||
return atoi(t)
|
||
case json.Number:
|
||
if n, err := t.Int64(); err == nil {
|
||
return int(n)
|
||
}
|
||
}
|
||
return def
|
||
}
|