1087 lines
28 KiB
Go
1087 lines
28 KiB
Go
package upstream
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"html"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"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
|
|
}
|
|
|
|
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
|
|
|
|
stateMu sync.Mutex
|
|
chatCfg map[string]chatConfig
|
|
cfgLocks map[string]*sync.Mutex
|
|
gates map[string]chan struct{}
|
|
|
|
// 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
|
|
|
|
// 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) *Client {
|
|
if log == nil {
|
|
log = slog.Default()
|
|
}
|
|
return &Client{
|
|
cfg: cfg,
|
|
ts: ts,
|
|
log: log,
|
|
chatCfg: map[string]chatConfig{},
|
|
cfgLocks: map[string]*sync.Mutex{},
|
|
gates: map[string]chan struct{}{},
|
|
quotaCool: map[string]time.Time{},
|
|
}
|
|
}
|
|
|
|
// 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()
|
|
until, ok := c.quotaCool[code]
|
|
now := time.Now()
|
|
if ok && now.Add(quotaCoolTTL).Before(until) {
|
|
c.stateMu.Unlock()
|
|
return
|
|
}
|
|
c.quotaCool[code] = now.Add(quotaCoolTTL)
|
|
c.stateMu.Unlock()
|
|
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))
|
|
}
|
|
|
|
// 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.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
|
|
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)
|
|
}
|
|
failed = append(failed, route)
|
|
c.logUpstreamErr(site.Code, modelID, attempt, route, err)
|
|
if isQuota(err) {
|
|
sawQuota = true
|
|
}
|
|
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.
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
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))
|
|
}
|
|
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 := map[string]any{}
|
|
_ = json.Unmarshal(resp.Body, &payload)
|
|
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, _, 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 AsUpstream(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
|
|
})
|
|
}
|
|
|
|
// 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 ────────────────────────────────────────────
|
|
|
|
func (c *Client) acquire(ctx context.Context, site string) error {
|
|
c.cfg.RLock()
|
|
max := c.cfg.Limits.MaxConcurrentPerSite
|
|
c.cfg.RUnlock()
|
|
if max <= 0 {
|
|
return nil
|
|
}
|
|
c.stateMu.Lock()
|
|
g, ok := c.gates[site]
|
|
if !ok {
|
|
g = make(chan struct{}, max)
|
|
c.gates[site] = g
|
|
}
|
|
c.stateMu.Unlock()
|
|
select {
|
|
case g <- struct{}{}:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return fmt.Errorf("site %s busy: %w", site, ctx.Err())
|
|
}
|
|
}
|
|
|
|
func (c *Client) release(site string) {
|
|
c.stateMu.Lock()
|
|
g, ok := c.gates[site]
|
|
c.stateMu.Unlock()
|
|
if !ok {
|
|
return
|
|
}
|
|
select {
|
|
case <-g:
|
|
default:
|
|
}
|
|
}
|
|
|
|
// ── 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
|
|
}
|
|
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
|
|
}
|