new-api-custom/model/detection.go

package model

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"time"

	"github.com/QuantumNous/new-api/common"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// DetectionRun is both the durable task and wallet ledger. Public projections
// are built explicitly; this record must never be serialized to anonymous users.
type DetectionRun struct {
	ID           string `gorm:"primaryKey;type:varchar(64)"`
	UserID       int    `gorm:"index"`
	RequestKey   string `gorm:"uniqueIndex;type:varchar(128)"`
	ChannelID    int    `gorm:"index"`
	GroupName    string
	Kind         string `gorm:"type:varchar(24)"`
	Status       string `gorm:"type:varchar(24);index"`
	CreatedAt    int64  `gorm:"autoCreateTime:false;index"`
	FinishedAt   int64
	Deadline     int64
	Reserved     int
	Charged      int
	Price        string `gorm:"type:text"`
	RemoteID     string
	Quality      string
	HTML         string `gorm:"type:text"`
	Result       string `gorm:"type:text"`
	Error        string
	Settled      bool
	LogJSON      string `gorm:"type:text"`
	LogDelivered bool   `gorm:"index"`
}

type DetectionCall struct {
	ID         string `gorm:"primaryKey;type:varchar(64)"`
	RunID      string `gorm:"index;type:varchar(64)"`
	SecretHash string `gorm:"uniqueIndex;type:varchar(64)"`
	ExpiresAt  int64
	StartedAt  int64
	FinishedAt int64
	State      string `gorm:"type:varchar(24)"`
	Usage      string `gorm:"type:text"`
	Sample     string `gorm:"type:text"`
	Quota      int
	Available  *bool
	Error      string
}

func DetectionSecretHash(secret string) string {
	sum := sha256.Sum256([]byte(secret))
	return hex.EncodeToString(sum[:])
}

func CreateDetection(run *DetectionRun) error {
	if run.Reserved < 0 || run.Reserved > common.MaxWalletQuota {
		return errors.New("invalid reservation")
	}
	original := *run
	err := relayAttemptTransaction(func(tx *gorm.DB) error {
		*run = original
		if run.UserID > 0 {
			var user User
			if err := lockForUpdate(tx).First(&user, run.UserID).Error; err != nil {
				return err
			}
			var previous DetectionRun
			if err := tx.Where("request_key = ?", run.RequestKey).First(&previous).Error; err == nil {
				if previous.UserID != run.UserID || previous.GroupName != run.GroupName || previous.Kind != run.Kind {
					return errors.New("idempotency key conflict")
				}
				*run = previous
				return nil
			} else if !errors.Is(err, gorm.ErrRecordNotFound) {
				return err
			}
			var active int64
			if err := tx.Model(&DetectionRun{}).Where("user_id = ? AND settled = ?", run.UserID, false).Count(&active).Error; err != nil {
				return err
			}
			if active > 0 {
				return errors.New("a detection is already running")
			}
			if err := tx.Model(&DetectionRun{}).Where("user_id = ? AND group_name = ? AND kind = ? AND created_at > ?", run.UserID, run.GroupName, run.Kind, time.Now().Unix()-300).Count(&active).Error; err != nil {
				return err
			}
			if active > 0 {
				return errors.New("please wait five minutes between detections")
			}
			if user.Status != common.UserStatusEnabled || user.Quota < run.Reserved {
				return errors.New("insufficient wallet balance")
			}
			if err := tx.Model(&User{}).Where("id = ?", run.UserID).Update("quota", gorm.Expr("quota - ?", run.Reserved)).Error; err != nil {
				return err
			}
		}
		return tx.Create(run).Error
	})
	if err == nil && run.UserID > 0 {
		_ = invalidateUserCache(run.UserID)
	}
	return err
}

func ClaimDetectionCall(hash string) (*DetectionCall, error) {
	var call DetectionCall
	err := relayAttemptTransaction(func(tx *gorm.DB) error {
		if err := tx.Where("secret_hash = ?", hash).First(&call).Error; err != nil {
			return err
		}
		var run DetectionRun
		if err := lockForUpdate(tx).First(&run, "id = ?", call.RunID).Error; err != nil {
			return err
		}
		if err := lockForUpdate(tx).Where("id = ?", call.ID).First(&call).Error; err != nil {
			return err
		}
		now := time.Now().Unix()
		if run.Kind != "pelican" || call.State != "ready" || call.ExpiresAt <= now || run.Settled || run.Status == "failed" || run.Status == "succeeded" || run.Deadline-10 <= now {
			return errors.New("expired or consumed credential")
		}
		update := tx.Model(&DetectionCall{}).Where("id = ? AND state = ?", call.ID, "ready").Updates(map[string]any{"state": "running", "started_at": now})
		if update.Error != nil {
			return update.Error
		}
		if update.RowsAffected != 1 {
			return errors.New("credential already consumed")
		}
		call.State = "running"
		call.StartedAt = now
		return nil
	})
	return &call, err
}

func FinishDetectionCall(call *DetectionCall) error {
	if call.Quota < 0 || call.Quota > common.MaxWalletQuota {
		return errors.New("invalid charge")
	}
	return DB.Model(&DetectionCall{}).Where("id = ? AND state = ?", call.ID, "running").Updates(map[string]any{
		"state": "finished", "finished_at": time.Now().Unix(), "usage": call.Usage, "sample": call.Sample, "quota": call.Quota, "available": call.Available, "error": call.Error,
	}).Error
}

func SettleDetection(id string) error {
	var owner int
	err := relayAttemptTransaction(func(tx *gorm.DB) error {
		var run DetectionRun
		if err := lockForUpdate(tx).First(&run, "id = ?", id).Error; err != nil {
			return err
		}
		owner = run.UserID
		if run.Settled {
			return nil
		}
		var calls []DetectionCall
		if err := tx.Where("run_id = ?", id).Find(&calls).Error; err != nil {
			return err
		}
		charge := 0
		for _, call := range calls {
			if call.State == "running" {
				return errors.New("generation still running")
			}
			if call.Quota < 0 || call.Quota > run.Reserved-charge {
				return errors.New("charge exceeds reservation")
			}
			charge += call.Quota
		}
		logJSON := ""
		if owner > 0 {
			var user User
			if err := lockForUpdate(tx).Unscoped().First(&user, owner).Error; err != nil {
				return err
			}
			if user.Quota > common.MaxWalletQuota-(run.Reserved-charge) || user.UsedQuota > common.MaxWalletQuota-charge {
				return errors.New("wallet settlement exceeds quota limit")
			}
			if err := tx.Unscoped().Model(&User{}).Where("id = ?", owner).Updates(map[string]any{"quota": gorm.Expr("quota + ?", run.Reserved-charge), "used_quota": gorm.Expr("used_quota + ?", charge)}).Error; err != nil {
				return err
			}
			key := "detection:" + id
			other, _ := common.Marshal(map[string]any{"billing_source": "wallet", "detection_id": id, "detection_kind": run.Kind, "reserved": run.Reserved, "refunded": run.Reserved - charge})
			entry := Log{AttemptKey: &key, UserId: owner, Username: user.Username, CreatedAt: time.Now().Unix(), Type: LogTypeSystem, Content: fmt.Sprintf("%s detection: charged %d quota; refunded %d quota", run.Kind, charge, run.Reserved-charge), Quota: charge, Group: run.GroupName, ModelName: "gpt-6-astra", RequestId: id, Other: string(other)}
			encoded, err := common.Marshal(entry)
			if err != nil {
				return err
			}
			logJSON = string(encoded)
		} else {
			charge = 0
		}
		if err := tx.Model(&DetectionCall{}).Where("run_id = ? AND state = ?", id, "ready").Update("state", "revoked").Error; err != nil {
			return err
		}
		return tx.Model(&DetectionRun{}).Where("id = ?", id).Updates(map[string]any{"settled": true, "charged": charge, "finished_at": time.Now().Unix(), "log_json": logJSON}).Error
	})
	if err == nil && owner > 0 {
		_ = invalidateUserCache(owner)
	}
	return err
}

// Detection ledger delivery is idempotent and separate from request metrics.
func DeliverDetectionLogs() error {
	var runs []DetectionRun
	if err := DB.Where("settled = ? AND log_delivered = ? AND log_json <> ?", true, false, "").Limit(100).Find(&runs).Error; err != nil {
		return err
	}
	for _, run := range runs {
		var entry Log
		if err := common.UnmarshalJsonStr(run.LogJSON, &entry); err != nil {
			return err
		}
		if err := LOG_DB.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "attempt_key"}}, DoNothing: true}).Create(&entry).Error; err != nil {
			return err
		}
		if err := DB.Model(&DetectionRun{}).Where("id = ?", run.ID).Update("log_delivered", true).Error; err != nil {
			return err
		}
	}
	return nil
}