package controller import ( "context" "crypto/sha256" "crypto/subtle" "encoding/hex" "errors" "net/http" "os" "slices" "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/gin-gonic/gin" "gorm.io/gorm" ) var chainEnqueue sync.Mutex var chainKinds = []string{"fingerprint", "ciallo", "vision"} // All three child tasks are committed together, using the existing durable ledger. func enqueueChainTest(ctx context.Context, group string, channelID int) ([]model.DetectionRun, error) { chainEnqueue.Lock() defer chainEnqueue.Unlock() var runs []model.DetectionRun err := model.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var active []model.DetectionRun if err := tx.Where("group_name = ? AND request_key LIKE ? AND settled = ? AND deadline > ?", group, "chain:%", false, time.Now().Unix()).Order("created_at asc").Limit(1).Find(&active).Error; err != nil { return err } if len(active) > 0 { id, _, _ := strings.Cut(active[0].ID, "_") return tx.Where("id IN ?", []string{id + "_fingerprint", id + "_ciallo", id + "_vision"}).Find(&runs).Error } id, err := common.GenerateRandomCharsKey(32) if err != nil { return err } now := time.Now().Unix() for _, kind := range chainKinds { runs = append(runs, model.DetectionRun{ID: id + "_" + kind, RequestKey: "chain:" + id + ":" + kind, ChannelID: channelID, GroupName: group, Kind: kind, Status: "pending", CreatedAt: now, Deadline: now + 600}) } return tx.Create(&runs).Error }) return runs, err } func BotChainTest(c *gin.Context) { c.Header("Cache-Control", "no-store") key := os.Getenv("BOT_BALANCE_KEY") derived := sha256.Sum256([]byte("newapi-chaintest-v1:" + key)) raw, bearer := strings.CutPrefix(c.GetHeader("Authorization"), "Bearer ") expected, supplied := sha256.Sum256([]byte(hex.EncodeToString(derived[:]))), sha256.Sum256([]byte(raw)) if len(key) < 32 || !bearer || subtle.ConstantTimeCompare(expected[:], supplied[:]) != 1 { c.JSON(http.StatusUnauthorized, gin.H{"success": false, "code": "unauthorized"}) return } var runs []model.DetectionRun if c.Request.Method == http.MethodPost { c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 1024) var input struct { Group string `json:"group"` } if c.ShouldBindJSON(&input) != nil || strings.TrimSpace(input.Group) == "" || len(input.Group) > 128 { c.JSON(400, gin.H{"success": false, "code": "invalid_group"}) return } input.Group = strings.TrimSpace(input.Group) channels, err := service.AstraSamplingChannels(c.Request.Context()) if err != nil { c.JSON(503, gin.H{"success": false, "code": "unavailable"}) return } slices.SortFunc(channels, func(a, b model.Channel) int { return a.Id - b.Id }) selected := 0 for _, channel := range channels { if slices.Contains(channel.GetGroups(), input.Group) { selected = channel.Id break } } if selected == 0 { c.JSON(404, gin.H{"success": false, "code": "group_unavailable"}) return } runs, err = enqueueChainTest(c.Request.Context(), input.Group, selected) if err != nil { c.JSON(503, gin.H{"success": false, "code": "unavailable"}) return } for _, run := range runs { go processDetection(context.Background(), run) } } else { id := c.Param("id") if len(id) != 32 || strings.ContainsAny(id, "_%") { c.JSON(404, gin.H{"success": false, "code": "not_found"}) return } err := model.DB.WithContext(c.Request.Context()).Where("id IN ? AND request_key LIKE ?", []string{id + "_fingerprint", id + "_ciallo", id + "_vision"}, "chain:"+id+":%").Find(&runs).Error if err != nil { c.JSON(503, gin.H{"success": false, "code": "unavailable"}) return } } data, err := chainTestResult(runs, time.Now().Unix()) if err != nil { c.JSON(404, gin.H{"success": false, "code": "not_found"}) return } c.JSON(200, gin.H{"success": true, "data": data}) } func chainTestResult(runs []model.DetectionRun, now int64) (gin.H, error) { if len(runs) != 3 { return nil, errors.New("incomplete chain") } id, _, _ := strings.Cut(runs[0].ID, "_") results := make(map[string]any, 3) done := true for _, run := range runs { run = service.NormalizeVisionRun(run) state, text, match := "running", "", "" if run.Settled { state = "failed" var result fingerprint.Result if run.Status == "succeeded" && common.UnmarshalJsonStr(run.Result, &result) == nil { text = result.Closest if run.Kind == "fingerprint" { match = service.AstraSamplingResultStatus(result) state = match if state != "match" && state != "mismatch" && state != "inconclusive" { state = "failed" } } else if strings.TrimSpace(text) != "" { state = service.VisionQuality(text) if run.Kind == "ciallo" { state = service.CialloQuality(text) } } } } else if now >= run.Deadline { state = "failed" } else { done = false } results[run.Kind] = gin.H{"status": state, "text": text} } return gin.H{"id": id, "group": runs[0].GroupName, "deadline": runs[0].Deadline, "done": done, "results": results}, nil }