cneicy
ceddfd3d6b
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
411 lines
13 KiB
Go
411 lines
13 KiB
Go
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/hmac"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
type detectionHandler struct{}
|
|
|
|
func pelicanCredential(id string) string {
|
|
mac := hmac.New(sha256.New, []byte(common.SessionSecret))
|
|
mac.Write([]byte("newapi-pelican-credential-v1:" + id))
|
|
return "det_" + hex.EncodeToString(mac.Sum(nil))
|
|
}
|
|
|
|
func (detectionHandler) Type() string { return "group_detection" }
|
|
func (detectionHandler) Enabled() bool {
|
|
return common.GetEnvOrDefaultBool("PELICAN_DETECTION_ENABLED", true)
|
|
}
|
|
func (detectionHandler) Interval() time.Duration { return 15 * time.Second }
|
|
func (detectionHandler) NewPayload() any { return nil }
|
|
|
|
func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner string) {
|
|
_ = model.DeliverDetectionLogs()
|
|
// Keep accounting facts, but expire bulky generated artwork after the public window.
|
|
_ = model.DB.Model(&model.DetectionRun{}).Where("settled = ? AND created_at < ? AND html <> ?", true, time.Now().Add(-6*time.Hour).Unix(), "").Update("html", "").Error
|
|
scheduleErr := schedulePelicanDetections(ctx)
|
|
var err error
|
|
{
|
|
var runs []model.DetectionRun
|
|
err = model.DB.Where("settled = ?", false).Order("created_at asc").Find(&runs).Error
|
|
if err == nil {
|
|
jobs := make(chan model.DetectionRun)
|
|
var workers sync.WaitGroup
|
|
for range 2 {
|
|
workers.Go(func() {
|
|
for run := range jobs {
|
|
processDetection(ctx, run)
|
|
}
|
|
})
|
|
}
|
|
for _, run := range runs {
|
|
select {
|
|
case jobs <- run:
|
|
case <-ctx.Done():
|
|
}
|
|
}
|
|
close(jobs)
|
|
workers.Wait()
|
|
}
|
|
}
|
|
if err == nil {
|
|
err = scheduleErr
|
|
}
|
|
status := model.SystemTaskStatusSucceeded
|
|
if err != nil {
|
|
status = model.SystemTaskStatusFailed
|
|
}
|
|
finishSystemTaskHandler(task, runner, status, nil, err)
|
|
}
|
|
|
|
func schedulePelicanDetections(ctx context.Context) error {
|
|
channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, "pelican", time.Now())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
now := time.Now().Unix()
|
|
for _, channel := range channels {
|
|
groups := channel.GetGroups()
|
|
if len(groups) == 0 {
|
|
continue
|
|
}
|
|
price, reserve, priceErr := service.DetectionPricing("", groups[0], "pelican")
|
|
if priceErr != nil {
|
|
return priceErr
|
|
}
|
|
created := time.Unix(now, 0)
|
|
price.Request.At = &created
|
|
encoded, _ := common.Marshal(price)
|
|
id, err := common.GenerateRandomCharsKey(32)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
run := &model.DetectionRun{ID: id, RequestKey: fmt.Sprintf("adaptive:%d:%s", channel.Id, schedules[channel.Id].LastResultID), ChannelID: channel.Id, GroupName: groups[0], Kind: "pelican", Status: "pending", CreatedAt: now, Price: string(encoded), Reserved: reserve}
|
|
if err = model.CreateDetection(run); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func processDetection(ctx context.Context, run model.DetectionRun) {
|
|
if ctx.Err() != nil {
|
|
return
|
|
}
|
|
if run.Deadline == 0 {
|
|
run.Deadline = time.Now().Unix() + 600
|
|
if model.DB.Model(&model.DetectionRun{}).Where("id = ? AND deadline = 0", run.ID).Updates(map[string]any{"deadline": run.Deadline, "status": "running"}).Error != nil {
|
|
return
|
|
}
|
|
}
|
|
taskCtx, taskCancel := context.WithDeadline(ctx, time.Unix(run.Deadline, 0))
|
|
defer taskCancel()
|
|
if time.Now().Unix() >= run.Deadline {
|
|
var calls []model.DetectionCall
|
|
if model.DB.Where("run_id = ?", run.ID).Find(&calls).Error != nil {
|
|
return
|
|
}
|
|
for _, call := range calls {
|
|
if call.State != "running" {
|
|
continue
|
|
}
|
|
if cancel, ok := detectionCalls.Load(call.ID); ok {
|
|
cancel.(context.CancelFunc)()
|
|
return
|
|
}
|
|
// All proxy contexts expire before the task deadline. A vanished process
|
|
// cannot resume or recharge this consumed credential.
|
|
call.Error = "interrupted_usage_unknown"
|
|
call.Quota = 0
|
|
_ = model.FinishDetectionCall(&call)
|
|
}
|
|
if run.RemoteID != "" {
|
|
cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
|
_, _ = manxueRunRequest(cancelCtx, &run, http.MethodDelete, "/api/v1/tests/"+run.RemoteID, nil, "")
|
|
cancel()
|
|
}
|
|
finishDetection(run, "failed", "detection_timeout")
|
|
return
|
|
}
|
|
if run.Status == "succeeded" || run.Status == "failed" {
|
|
settleDetectionAndSchedule(run)
|
|
return
|
|
}
|
|
if run.Kind == "fingerprint" {
|
|
processPaidFingerprint(ctx, run)
|
|
return
|
|
}
|
|
if run.RemoteID == localPelicanID {
|
|
processLocalPelican(taskCtx, run)
|
|
return
|
|
}
|
|
if run.RemoteID == "" {
|
|
var existing int64
|
|
if model.DB.Model(&model.DetectionCall{}).Where("run_id = ?", run.ID).Count(&existing).Error != nil {
|
|
return
|
|
}
|
|
// A domain-separated PRF allows safe retry after a submit/commit crash;
|
|
// only its digest is stored, and the random run ID prevents enumeration.
|
|
secret := pelicanCredential(run.ID)
|
|
call := model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(secret), ExpiresAt: run.Deadline + 300, State: "ready"}
|
|
if existing == 0 && model.DB.Create(&call).Error != nil {
|
|
return
|
|
}
|
|
payload := map[string]any{"benchmark": "pelican", "base_url": "https://ai.shrinken.pw/api/detection-proxy/v1", "api_key": secret, "model": fingerprint.Model, "protocol": "responses", "reasoning_effort": "high"}
|
|
var result gjson.Result
|
|
var err error
|
|
for taskCtx.Err() == nil {
|
|
result, err = manxueRunRequest(taskCtx, &run, http.MethodPost, "/api/v1/tests", payload, run.ID)
|
|
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
|
|
processLocalPelican(taskCtx, run)
|
|
return
|
|
}
|
|
if err == nil {
|
|
break
|
|
}
|
|
select {
|
|
case <-taskCtx.Done():
|
|
case <-time.After(5 * time.Second):
|
|
}
|
|
}
|
|
if err != nil || taskCtx.Err() != nil {
|
|
if ctx.Err() == nil {
|
|
processDetection(ctx, run)
|
|
}
|
|
return
|
|
}
|
|
remoteID := result.Get("id").String()
|
|
if remoteID == "" || len(remoteID) > 128 || strings.ContainsAny(remoteID, "/?#\\") {
|
|
finishDetection(run, "failed", "invalid_detection_response")
|
|
return
|
|
}
|
|
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"remote_id": remoteID, "status": "running"}).Error != nil {
|
|
return
|
|
}
|
|
run.RemoteID = remoteID
|
|
}
|
|
// Each worker retains its slot until completion, limiting remote generation
|
|
// concurrency as well as local submission concurrency.
|
|
for time.Now().Unix() < run.Deadline {
|
|
result, err := manxueRunRequest(taskCtx, &run, http.MethodGet, "/api/v1/tests/"+run.RemoteID, nil, "")
|
|
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
|
|
processLocalPelican(taskCtx, run)
|
|
return
|
|
}
|
|
if err != nil {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-time.After(5 * time.Second):
|
|
}
|
|
continue
|
|
}
|
|
switch result.Get("status").String() {
|
|
case "succeeded":
|
|
quality := result.Get("assessment.quality").String()
|
|
if quality != "normal" && quality != "degraded" {
|
|
quality = "unknown"
|
|
}
|
|
html := result.Get("result.html").String()
|
|
if len(html) > 2<<20 {
|
|
html = ""
|
|
}
|
|
metadata, _ := common.Marshal(map[string]string{"source_url": manxueArtworkURL(result)})
|
|
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": quality, "html": html, "result": string(metadata)}).Error != nil {
|
|
return
|
|
}
|
|
finishDetection(run, "succeeded", "")
|
|
return
|
|
case "failed", "cancelled":
|
|
finishDetection(run, "failed", "detection_service_failed")
|
|
return
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-time.After(5 * time.Second):
|
|
}
|
|
}
|
|
processDetection(ctx, run)
|
|
}
|
|
|
|
func finishDetection(run model.DetectionRun, status, reason string) {
|
|
updates := map[string]any{"status": status, "error": reason}
|
|
if status == "failed" && run.Kind == "pelican" {
|
|
updates["result"] = ""
|
|
}
|
|
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(updates).Error != nil {
|
|
return
|
|
}
|
|
settleDetectionAndSchedule(run)
|
|
}
|
|
|
|
func settleDetectionAndSchedule(run model.DetectionRun) {
|
|
if model.SettleDetection(run.ID) != nil || run.UserID != 0 || run.Kind != "pelican" {
|
|
return
|
|
}
|
|
// Reconciliation also repairs the clock if a process exits after settlement.
|
|
if err := model.DB.Select("id", "channel_id", "finished_at", "status", "quality").First(&run, "id = ?", run.ID).Error; err != nil {
|
|
return
|
|
}
|
|
outcome := "unknown"
|
|
if run.Status == "succeeded" {
|
|
outcome = run.Quality
|
|
}
|
|
_ = model.ApplyDetectionScheduleResult(context.Background(), run.ChannelID, "pelican", run.ID, run.FinishedAt, outcome)
|
|
}
|
|
|
|
func manxueRequest(ctx context.Context, method, path string, payload any, key string) (gjson.Result, error) {
|
|
client := &http.Client{Timeout: 35 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
|
return manxueRequestWithClient(ctx, client, method, path, payload, key)
|
|
}
|
|
|
|
func manxueRequestWithClient(ctx context.Context, client *http.Client, method, path string, payload any, key string) (gjson.Result, error) {
|
|
var encoded []byte
|
|
if payload != nil {
|
|
var err error
|
|
encoded, err = common.Marshal(payload)
|
|
if err != nil {
|
|
return gjson.Result{}, err
|
|
}
|
|
}
|
|
for attempt := 0; attempt < 3; attempt++ {
|
|
req, err := http.NewRequestWithContext(ctx, method, "https://manxue.ai"+path, bytes.NewReader(encoded))
|
|
if err != nil {
|
|
return gjson.Result{}, err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("User-Agent", "Mozilla/5.0")
|
|
if strings.HasPrefix(path, "/api/visitor/") {
|
|
req.Header.Set("Origin", "https://manxue.ai")
|
|
}
|
|
if key != "" {
|
|
req.Header.Set("Idempotency-Key", key)
|
|
}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return gjson.Result{}, errDetectionServiceUnavailable
|
|
}
|
|
body, readErr := io.ReadAll(io.LimitReader(resp.Body, (3<<20)+1))
|
|
resp.Body.Close()
|
|
if readErr != nil || resp.StatusCode >= 500 {
|
|
return gjson.Result{}, errDetectionServiceUnavailable
|
|
}
|
|
if resp.StatusCode == 429 || resp.StatusCode == 503 {
|
|
seconds, _ := strconv.Atoi(resp.Header.Get("Retry-After"))
|
|
if at, err := http.ParseTime(resp.Header.Get("Retry-After")); err == nil {
|
|
seconds = int(time.Until(at).Seconds()) + 1
|
|
}
|
|
seconds = max(seconds, 2<<attempt)
|
|
select {
|
|
case <-ctx.Done():
|
|
return gjson.Result{}, ctx.Err()
|
|
case <-time.After(time.Duration(seconds) * time.Second):
|
|
}
|
|
continue
|
|
}
|
|
if readErr != nil || len(body) > 3<<20 || resp.StatusCode < 200 || resp.StatusCode >= 300 || !gjson.ValidBytes(body) {
|
|
return gjson.Result{}, errors.New("invalid detection service response")
|
|
}
|
|
return gjson.ParseBytes(body), nil
|
|
}
|
|
return gjson.Result{}, errDetectionServiceUnavailable
|
|
}
|
|
|
|
func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
|
|
channel, err := model.GetChannelById(run.ChannelID, true)
|
|
if err != nil || channel.Status != common.ChannelStatusEnabled {
|
|
finishDetection(run, "failed", "channel_unavailable")
|
|
return
|
|
}
|
|
samples := []fingerprint.Sample{}
|
|
for index := 0; index < 3; index++ {
|
|
id := fmt.Sprintf("%s_%d", run.ID, index)
|
|
var previous []model.DetectionCall
|
|
if model.DB.Where("id = ?", id).Find(&previous).Error != nil {
|
|
return
|
|
}
|
|
if len(previous) > 0 {
|
|
if previous[0].State == "running" {
|
|
return
|
|
}
|
|
var sample fingerprint.Sample
|
|
if previous[0].Sample != "" && common.UnmarshalJsonStr(previous[0].Sample, &sample) == nil {
|
|
samples = append(samples, sample)
|
|
}
|
|
continue
|
|
}
|
|
call := model.DetectionCall{ID: id, RunID: run.ID, SecretHash: model.DetectionSecretHash(id), State: "running", StartedAt: time.Now().Unix()}
|
|
if model.DB.Create(&call).Error != nil {
|
|
return
|
|
}
|
|
probeCtx, cancel := context.WithDeadline(ctx, time.Unix(run.Deadline-10, 0))
|
|
detectionCalls.Store(call.ID, cancel)
|
|
answer, count, probeErr := runAstraFingerprintProbe(probeCtx, channel, true)
|
|
interrupted := probeCtx.Err() != nil
|
|
cancel()
|
|
available := probeErr == nil || answer.HasUsage
|
|
call.Available = &available
|
|
if !answer.HasUsage && (errors.Is(probeErr, fingerprint.ErrInvalidSample) || errors.Is(probeErr, fingerprint.ErrIncomplete) || interrupted) {
|
|
call.Available = nil
|
|
}
|
|
if probeErr != nil {
|
|
call.Error = "channel_call_failed"
|
|
}
|
|
if answer.HasUsage {
|
|
usage, _ := common.Marshal(answer)
|
|
call.Usage = string(usage)
|
|
quota, quotaErr := service.DetectionUsageQuota(run.Price, answer.InputTokens, answer.OutputTokens, answer.CachedTokens, answer.CacheWriteTokens)
|
|
if quotaErr != nil || quota > run.Reserved/3 {
|
|
call.Error = "settlement_unavailable"
|
|
} else {
|
|
call.Quota = quota
|
|
}
|
|
} else {
|
|
call.Error = "usage_unavailable"
|
|
}
|
|
if probeErr == nil && fingerprint.ValidArray(answer.Text, count) {
|
|
sample := fingerprint.Sample{Text: answer.Text, Count: count}
|
|
samples = append(samples, sample)
|
|
encoded, _ := common.Marshal(sample)
|
|
call.Sample = string(encoded)
|
|
}
|
|
if model.FinishDetectionCall(&call) != nil {
|
|
detectionCalls.Delete(call.ID)
|
|
return
|
|
}
|
|
detectionCalls.Delete(call.ID)
|
|
if ctx.Err() != nil {
|
|
return
|
|
}
|
|
}
|
|
result, err := fingerprint.Score(samples)
|
|
if err != nil {
|
|
finishDetection(run, "failed", "scoring_failed")
|
|
return
|
|
}
|
|
encoded, _ := common.Marshal(result)
|
|
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": result.Status, "result": string(encoded)}).Error != nil {
|
|
return
|
|
}
|
|
finishDetection(run, "succeeded", "")
|
|
}
|