senlinai/agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden controlled AI session consistency

Browse Source
This commit is contained in:
yanweidong
2026-07-21 19:54:25 +08:00
parent 122f5d8c52
commit 251b212e06
15 changed files with 406 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
backend/internal/logic/ai/gateway.go
View File

@@ -18,6 +18,7 @@ import (
type Gateway struct {
systemKey string
encryptionSecret string
now func() time.Time
}
type SelectedKey struct {
@@ -36,7 +37,7 @@ func NewGateway(systemKey string) *Gateway {
}
func NewGatewayWithSecret(systemKey string, encryptionSecret string) *Gateway {
return &Gateway{systemKey: systemKey, encryptionSecret: encryptionSecret}
return &Gateway{systemKey: systemKey, encryptionSecret: encryptionSecret, now: time.Now}
}
func (g *Gateway) SaveUserKey(userID uint, provider string, apiKey string) error {
@@ -55,11 +56,13 @@ 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 SelectedKey{}, err
return selected, err
}
return SelectedKey{Provider: userKey.Provider, APIKey: apiKey, KeyType: "user"}, nil
selected.APIKey = apiKey
return selected, nil
}
if !errors.Is(err, gorm.ErrRecordNotFound) {
return SelectedKey{}, err
@@ -70,8 +73,11 @@ func (g *Gateway) SelectKey(userID uint) (SelectedKey, error) {
return SelectedKey{Provider: "openai", APIKey: g.systemKey, KeyType: "system"}, nil
}
func (g *Gateway) RecordCall(userID uint, provider string, usedKeyType string, action string, status string, errText string) error {
return models.DBService.Create(&models.SenlinAgentAICallLog{
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,
@@ -81,17 +87,34 @@ func (g *Gateway) RecordCall(userID uint, provider string, usedKeyType string, a
}).Error
}
func (g *Gateway) CheckRateLimit(userID uint, action string, limit int, window time.Duration) 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
}
var count int64
if err := models.DBService.Model(&models.SenlinAgentAICallLog{}).
Where("user_id = ? AND action = ? AND created_at >= ?", userID, action, time.Now().Add(-window)).
Count(&count).Error; err != nil {
return err
currentTime := time.Now().UTC()
if g.now != nil {
currentTime = g.now().UTC()
}
if count >= int64(limit) {
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