231 lines
7.9 KiB
Go
231 lines
7.9 KiB
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 {
|
|
FirstTokenMS *int64
|
|
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, "first_token_ms": call.FirstTokenMS,
|
|
}).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
|
|
}
|