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

controller/detection_worker.go

@@ -1,22 +1,15 @@
 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"
+	benchmark "github.com/QuantumNous/new-api/pkg/model_benchmark"
 	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
 	"github.com/QuantumNous/new-api/service"
 	"github.com/tidwall/gjson"
@@ -24,15 +17,9 @@ import (
 
 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)
+	return common.GetEnvOrDefaultBool("MODEL_DETECTION_ENABLED", true)
 }
 func (detectionHandler) Interval() time.Duration { return 15 * time.Second }
 func (detectionHandler) NewPayload() any         { return nil }
@@ -41,7 +28,10 @@ func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner
 	_ = 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)
+	scheduleErr := scheduleBenchmarkDetections(ctx)
+	if candyErr := scheduleAutomaticDetections(ctx, "candy"); scheduleErr == nil {
+		scheduleErr = candyErr
+	}
 	var err error
 	{
 		var runs []model.DetectionRun
@@ -76,8 +66,12 @@ func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner
 	finishSystemTaskHandler(task, runner, status, nil, err)
 }
 
-func schedulePelicanDetections(ctx context.Context) error {
-	channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, "pelican", time.Now())
+func scheduleBenchmarkDetections(ctx context.Context) error {
+	return scheduleAutomaticDetections(ctx, "benchmark")
+}
+
+func scheduleAutomaticDetections(ctx context.Context, kind string) error {
+	channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, kind, time.Now())
 	if err != nil {
 		return err
 	}
@@ -87,7 +81,7 @@ func schedulePelicanDetections(ctx context.Context) error {
 		if len(groups) == 0 {
 			continue
 		}
-		price, reserve, priceErr := service.DetectionPricing("", groups[0], "pelican")
+		price, reserve, priceErr := service.DetectionPricing("", groups[0], kind)
 		if priceErr != nil {
 			return priceErr
 		}
@@ -98,7 +92,7 @@ func schedulePelicanDetections(ctx context.Context) error {
 		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}
+		run := &model.DetectionRun{ID: id, RequestKey: fmt.Sprintf("%s:%d:%s", kind, channel.Id, schedules[channel.Id].LastResultID), ChannelID: channel.Id, GroupName: groups[0], Kind: kind, Status: "pending", CreatedAt: now, Price: string(encoded), Reserved: reserve}
 		if err = model.CreateDetection(run); err != nil {
 			return err
 		}
@@ -116,8 +110,6 @@ func processDetection(ctx context.Context, run model.DetectionRun) {
 			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 {
@@ -137,11 +129,6 @@ func processDetection(ctx context.Context, run model.DetectionRun) {
 			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
 	}
@@ -149,107 +136,16 @@ func processDetection(ctx context.Context, run model.DetectionRun) {
 		settleDetectionAndSchedule(run)
 		return
 	}
-	if run.Kind == "fingerprint" {
-		processPaidFingerprint(ctx, run)
+	if run.Kind != "fingerprint" && run.Kind != "benchmark" && run.Kind != "candy" {
+		finishDetection(run, "failed", "detection_retired")
 		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)
+	processPaidFingerprint(ctx, run)
 }
 
 func finishDetection(run model.DetectionRun, status, reason string) {
 	updates := map[string]any{"status": status, "error": reason}
-	if status == "failed" && run.Kind == "pelican" {
+	if status == "failed" && run.Kind == "benchmark" {
 		updates["result"] = ""
 	}
 	if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(updates).Error != nil {
@@ -259,7 +155,7 @@ func finishDetection(run model.DetectionRun, status, reason string) {
 }
 
 func settleDetectionAndSchedule(run model.DetectionRun) {
-	if model.SettleDetection(run.ID) != nil || run.UserID != 0 || run.Kind != "pelican" {
+	if model.SettleDetection(run.ID) != nil || run.UserID != 0 || (run.Kind != "benchmark" && run.Kind != "candy") {
 		return
 	}
 	// Reconciliation also repairs the clock if a process exits after settlement.
@@ -270,64 +166,7 @@ func settleDetectionAndSchedule(run model.DetectionRun) {
 	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
+	_ = model.ApplyDetectionScheduleResult(context.Background(), run.ChannelID, run.Kind, run.ID, run.FinishedAt, outcome)
 }
 
 func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
@@ -336,8 +175,20 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
 		finishDetection(run, "failed", "channel_unavailable")
 		return
 	}
+	output, calls := fingerprint.OutputTokens, fingerprint.ProbeCount
+	if run.Kind == "benchmark" {
+		output, calls = benchmark.OutputTokens, benchmark.ProbeCount
+	}
+	if run.Kind == "candy" {
+		output, calls = service.CandyOutputTokens, 1
+	}
+	var price service.DetectionPrice
+	if common.UnmarshalJsonStr(run.Price, &price) != nil || gjson.GetBytes(price.Request.Body, "max_output_tokens").Int() != int64(output) {
+		finishDetection(run, "failed", "benchmark_changed")
+		return
+	}
 	samples := []fingerprint.Sample{}
-	for index := 0; index < 3; index++ {
+	for index := 0; index < calls; index++ {
 		id := fmt.Sprintf("%s_%d", run.ID, index)
 		var previous []model.DetectionCall
 		if model.DB.Where("id = ?", id).Find(&previous).Error != nil {
@@ -359,7 +210,7 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
 		}
 		probeCtx, cancel := context.WithDeadline(ctx, time.Unix(run.Deadline-10, 0))
 		detectionCalls.Store(call.ID, cancel)
-		answer, count, probeErr := runAstraFingerprintProbe(probeCtx, channel, true)
+		answer, count, probeErr := runModelDetectionProbe(probeCtx, channel, index, run.Kind)
 		interrupted := probeCtx.Err() != nil
 		cancel()
 		available := probeErr == nil || answer.HasUsage
@@ -374,7 +225,7 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
 			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 {
+			if quotaErr != nil || quota > run.Reserved/calls {
 				call.Error = "settlement_unavailable"
 			} else {
 				call.Quota = quota
@@ -382,7 +233,7 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
 		} else {
 			call.Error = "usage_unavailable"
 		}
-		if probeErr == nil && fingerprint.ValidArray(answer.Text, count) {
+		if probeErr == nil && ((run.Kind == "candy" && answer.Text != "") || (run.Kind == "benchmark" && benchmark.ValidSample(answer.Text)) || (run.Kind == "fingerprint" && fingerprint.ValidArray(answer.Text, count))) {
 			sample := fingerprint.Sample{Text: answer.Text, Count: count}
 			samples = append(samples, sample)
 			encoded, _ := common.Marshal(sample)
@@ -397,13 +248,34 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
 			return
 		}
 	}
-	result, err := fingerprint.Score(samples)
+	var result fingerprint.Result
+	if run.Kind == "benchmark" {
+		secondary := make([]benchmark.Sample, len(samples))
+		for i, sample := range samples {
+			secondary[i] = benchmark.Sample{Text: sample.Text, Count: sample.Count}
+		}
+		scored, scoreErr := benchmark.Score(secondary)
+		err = scoreErr
+		result = fingerprint.Result{Status: scored.Status, Engine: scored.Engine, ValidSamples: scored.ValidSamples, Closest: scored.Closest, Weight: scored.Weight, Gap: scored.Gap, Matches: scored.Matches}
+	} else if run.Kind == "candy" {
+		if len(samples) != 1 {
+			finishDetection(run, "failed", "answer_unavailable")
+			return
+		}
+		result = fingerprint.Result{Status: service.CandyQuality(samples[0].Text), Engine: "candy-v1", ValidSamples: 1, Closest: samples[0].Text}
+	} else {
+		result, err = fingerprint.Score(samples)
+	}
 	if err != nil {
 		finishDetection(run, "failed", "scoring_failed")
 		return
 	}
+	quality := result.Status
+	if run.Kind == "benchmark" {
+		quality = service.FingerprintScheduleOutcome(result.Status)
+	}
 	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 {
+	if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": quality, "result": string(encoded)}).Error != nil {
 		return
 	}
 	finishDetection(run, "succeeded", "")