new-api-custom/controller/detection_worker.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", "")
}