agent/backend/internal/logic/ai/gateway.go

package ai

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"io"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"senlinai-agent/backend/internal/models"
)

type Gateway struct {
	systemKey        string
	encryptionSecret string
	now              func() time.Time
}

type SelectedKey struct {
	Provider string
	APIKey   string
	KeyType  string
}

var (
	ErrAIKeyMissing  = errors.New("no ai key available")
	ErrAIRateLimited = errors.New("ai rate limit exceeded")
)

func NewGateway(systemKey string) *Gateway {
	return NewGatewayWithSecret(systemKey, "development-ai-key-secret-change-me")
}

func NewGatewayWithSecret(systemKey string, encryptionSecret string) *Gateway {
	return &Gateway{systemKey: systemKey, encryptionSecret: encryptionSecret, now: time.Now}
}

func (g *Gateway) SaveUserKey(userID uint, provider string, apiKey string) error {
	encrypted, err := encryptAPIKey(apiKey, g.encryptionSecret)
	if err != nil {
		return err
	}
	key := models.SenlinAgentAIKey{UserID: userID, Provider: provider, EncryptedAPIKey: encrypted}
	return models.DBService.Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "user_id"}},
		DoUpdates: clause.AssignmentColumns([]string{"provider", "encrypted_api_key", "updated_at"}),
	}).Create(&key).Error
}

func (g *Gateway) SelectKey(userID uint) (SelectedKey, error) {
	var userKey models.SenlinAgentAIKey
	err := models.DBService.Where("user_id = ?", userID).First(&userKey).Error
	if err == nil {
		selected := SelectedKey{Provider: userKey.Provider, KeyType: "user"}
		apiKey, err := decryptAPIKey(userKey.EncryptedAPIKey, g.encryptionSecret)
		if err != nil {
			return selected, err
		}
		selected.APIKey = apiKey
		return selected, nil
	}
	if !errors.Is(err, gorm.ErrRecordNotFound) {
		return SelectedKey{}, err
	}
	if g.systemKey == "" {
		return SelectedKey{}, ErrAIKeyMissing
	}
	return SelectedKey{Provider: "openai", APIKey: g.systemKey, KeyType: "system"}, nil
}

func (g *Gateway) RecordCall(database *gorm.DB, userID uint, provider string, usedKeyType string, action string, status string, errText string) error {
	if database == nil {
		database = models.DBService
	}
	return database.Create(&models.SenlinAgentAICallLog{
		UserID:      userID,
		Provider:    provider,
		UsedKeyType: usedKeyType,
		Action:      action,
		Status:      status,
		Error:       errText,
	}).Error
}

// ReserveRateLimit 以数据库单条 UPSERT 原子占用固定窗口配额。
// 配额在 provider/key 选择前占用，后续缺 key 或 provider 失败同样计入该窗口的尝试次数。
func (g *Gateway) ReserveRateLimit(userID uint, action string, limit int, window time.Duration) error {
	if limit <= 0 {
		return nil
	}
	currentTime := time.Now().UTC()
	if g.now != nil {
		currentTime = g.now().UTC()
	}
	windowStart := currentTime.Truncate(window)
	bucket := models.SenlinAgentAIRateBucket{
		UserID: userID, Action: action, WindowStart: windowStart, Count: 1,
	}
	result := models.DBService.Clauses(clause.OnConflict{
		Columns: []clause.Column{{Name: "user_id"}, {Name: "action"}, {Name: "window_start"}},
		DoUpdates: clause.Assignments(map[string]any{
			"count":      gorm.Expr("senlin_agent_ai_rate_buckets.count + 1"),
			"updated_at": currentTime,
		}),
		Where: clause.Where{Exprs: []clause.Expression{
			clause.Lt{Column: clause.Column{Table: "senlin_agent_ai_rate_buckets", Name: "count"}, Value: limit},
		}},
	}).Create(&bucket)
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return ErrAIRateLimited
	}
	return nil
}

func encryptAPIKey(apiKey string, secret string) (string, error) {
	if apiKey == "" {
		return "", errors.New("api key is required")
	}
	block, err := aes.NewCipher(encryptionKey(secret))
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}
	ciphertext := gcm.Seal(nil, nonce, []byte(apiKey), nil)
	return "v1:" + base64.RawURLEncoding.EncodeToString(append(nonce, ciphertext...)), nil
}

func decryptAPIKey(encrypted string, secret string) (string, error) {
	if len(encrypted) < 3 || encrypted[:3] != "v1:" {
		return encrypted, nil
	}
	raw, err := base64.RawURLEncoding.DecodeString(encrypted[3:])
	if err != nil {
		return "", errors.New("invalid encrypted api key")
	}
	block, err := aes.NewCipher(encryptionKey(secret))
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	if len(raw) < gcm.NonceSize() {
		return "", errors.New("invalid encrypted api key")
	}
	nonce := raw[:gcm.NonceSize()]
	ciphertext := raw[gcm.NonceSize():]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", errors.New("invalid encrypted api key")
	}
	return string(plaintext), nil
}

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