feat: sync deployed v1.0.0-rc.38-custom.42

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.

controller/detection_worker.go

@@ -0,0 +1,410 @@
+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", "")
+}