Skip to content

Commit 5fcdb1f

Browse files
Bowl42Bowl42
committed
feat: 添加 Antigravity 额度管理功能
- 新增后台任务服务 (AntigravityTaskService) 定期刷新配额 - 在 Provider 列表添加手动刷新额度按钮 - 显示 Image 模型额度 (alongside Claude quota) - 添加 /antigravity/refresh-quotas API 端点
1 parent 88e9d16 commit 5fcdb1f

File tree

16 files changed

+900
-100
lines changed

16 files changed

+900
-100
lines changed

‎cmd/maxx/main.go‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -174,16 +174,27 @@ func main() {
174174
}()
175175
log.Println("[Cooldown] Background cleanup started (runs every 1 hour)")
176176

177+
// Create WebSocket hub
178+
wsHub := handler.NewWebSocketHub()
179+
180+
// Create Antigravity task service for periodic quota refresh and auto-sorting
181+
antigravityTaskSvc := service.NewAntigravityTaskService(
182+
cachedProviderRepo,
183+
cachedRouteRepo,
184+
antigravityQuotaRepo,
185+
settingRepo,
186+
proxyRequestRepo,
187+
wsHub,
188+
)
189+
177190
// Start background tasks
178191
core.StartBackgroundTasks(core.BackgroundTaskDeps{
179-
UsageStats: usageStatsRepo,
180-
ProxyRequest: proxyRequestRepo,
181-
Settings: settingRepo,
192+
UsageStats: usageStatsRepo,
193+
ProxyRequest: proxyRequestRepo,
194+
Settings: settingRepo,
195+
AntigravityTaskSvc: antigravityTaskSvc,
182196
})
183197

184-
// Create WebSocket hub
185-
wsHub := handler.NewWebSocketHub()
186-
187198
// Setup log output to broadcast via WebSocket
188199
logWriter := handler.NewWebSocketLogWriter(wsHub, os.Stdout, logPath)
189200
log.SetOutput(logWriter)
@@ -251,6 +262,7 @@ func main() {
251262
adminHandler := handler.NewAdminHandler(adminService, backupService, logPath)
252263
authHandler := handler.NewAuthHandler(authMiddleware)
253264
antigravityHandler := handler.NewAntigravityHandler(adminService, antigravityQuotaRepo, wsHub)
265+
antigravityHandler.SetTaskService(antigravityTaskSvc)
254266
kiroHandler := handler.NewKiroHandler(adminService)
255267

256268
// Use already-created cached project repository for project proxy handler

‎internal/core/task.go‎

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
package core
22

33
import (
4+
"context"
45
"log"
56
"strconv"
67
"time"
78

89
"github.com/awsl-project/maxx/internal/domain"
910
"github.com/awsl-project/maxx/internal/repository"
11+
"github.com/awsl-project/maxx/internal/service"
1012
)
1113

1214
const (
@@ -15,9 +17,10 @@ const (
1517

1618
// BackgroundTaskDeps 后台任务依赖
1719
type BackgroundTaskDeps struct {
18-
UsageStats repository.UsageStatsRepository
19-
ProxyRequest repository.ProxyRequestRepository
20-
Settings repository.SystemSettingRepository
20+
UsageStats repository.UsageStatsRepository
21+
ProxyRequest repository.ProxyRequestRepository
22+
Settings repository.SystemSettingRepository
23+
AntigravityTaskSvc *service.AntigravityTaskService
2124
}
2225

2326
// StartBackgroundTasks 启动所有后台任务
@@ -66,6 +69,11 @@ func StartBackgroundTasks(deps BackgroundTaskDeps) {
6669
}
6770
}()
6871

72+
// Antigravity 配额刷新任务（动态间隔）
73+
if deps.AntigravityTaskSvc != nil {
74+
go deps.runAntigravityQuotaRefresh()
75+
}
76+
6977
log.Println("[Task] Background tasks started (minute:30s, hour:1m, day:5m, cleanup:1h)")
7078
}
7179

@@ -124,3 +132,24 @@ func (d *BackgroundTaskDeps) cleanupOldRequests() {
124132
log.Printf("[Task] Deleted %d requests older than %d hours", deleted, retentionHours)
125133
}
126134
}
135+
136+
// runAntigravityQuotaRefresh 定期刷新 Antigravity 配额
137+
func (d *BackgroundTaskDeps) runAntigravityQuotaRefresh() {
138+
time.Sleep(30 * time.Second) // 初始延迟
139+
140+
for {
141+
interval := d.AntigravityTaskSvc.GetRefreshInterval()
142+
if interval <= 0 {
143+
// 禁用状态，每分钟检查一次配置
144+
time.Sleep(1 * time.Minute)
145+
continue
146+
}
147+
148+
// 执行刷新
149+
ctx := context.Background()
150+
d.AntigravityTaskSvc.RefreshQuotas(ctx)
151+
152+
// 等待下一次刷新
153+
time.Sleep(time.Duration(interval) * time.Minute)
154+
}
155+
}

‎internal/domain/model.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -368,9 +368,11 @@ type SystemSetting struct {
368368

369369
// 系统设置 Key 常量
370370
const (
371-
SettingKeyProxyPort = "proxy_port" // 代理服务器端口，默认 9880
372-
SettingKeyRequestRetentionHours = "request_retention_hours" // 请求记录保留小时数，默认 168 小时（7天），0 表示不清理
373-
SettingKeyTimezone = "timezone" // 时区设置，默认 Asia/Shanghai
371+
SettingKeyProxyPort = "proxy_port" // 代理服务器端口，默认 9880
372+
SettingKeyRequestRetentionHours = "request_retention_hours" // 请求记录保留小时数，默认 168 小时（7天），0 表示不清理
373+
SettingKeyTimezone = "timezone" // 时区设置，默认 Asia/Shanghai
374+
SettingKeyQuotaRefreshInterval = "quota_refresh_interval" // Antigravity 配额刷新间隔（分钟），0 表示禁用
375+
SettingKeyAutoSortAntigravity = "auto_sort_antigravity" // 是否自动排序 Antigravity 路由，"true" 或 "false"
374376
)
375377

376378
// Antigravity 模型配额

‎internal/handler/antigravity.go‎

Lines changed: 54 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ type AntigravityHandler struct {
2121
svc *service.AdminService
2222
quotaRepo repository.AntigravityQuotaRepository
2323
oauthManager *antigravity.OAuthManager
24+
taskSvc *service.AntigravityTaskService
2425
}
2526

2627
// NewAntigravityHandler creates a new Antigravity handler
@@ -32,6 +33,11 @@ func NewAntigravityHandler(svc *service.AdminService, quotaRepo repository.Antig
3233
}
3334
}
3435

36+
// SetTaskService sets the AntigravityTaskService for background task operations
37+
func (h *AntigravityHandler) SetTaskService(taskSvc *service.AntigravityTaskService) {
38+
h.taskSvc = taskSvc
39+
}
40+
3541
// ServeHTTP routes Antigravity requests
3642
// Routes:
3743
// POST /antigravity/validate-token - 验证单个 refresh token
@@ -40,6 +46,8 @@ func NewAntigravityHandler(svc *service.AdminService, quotaRepo repository.Antig
4046
// GET /antigravity/providers/quotas - 批量获取所有 Antigravity provider 的配额信息
4147
// POST /antigravity/oauth/start - 启动 OAuth 流程
4248
// GET /antigravity/oauth/callback - OAuth 回调
49+
// POST /antigravity/refresh-quotas - 强制刷新所有配额
50+
// POST /antigravity/sort-routes - 手动排序路由
4351
func (h *AntigravityHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
4452
path := strings.TrimPrefix(r.URL.Path, "/antigravity")
4553
path = strings.TrimSuffix(path, "/")
@@ -58,6 +66,18 @@ func (h *AntigravityHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
5866
return
5967
}
6068

69+
// POST /antigravity/refresh-quotas - 强制刷新所有配额
70+
if len(parts) >= 2 && parts[1] == "refresh-quotas" && r.Method == http.MethodPost {
71+
h.handleForceRefreshQuotas(w, r)
72+
return
73+
}
74+
75+
// POST /antigravity/sort-routes - 手动排序路由
76+
if len(parts) >= 2 && parts[1] == "sort-routes" && r.Method == http.MethodPost {
77+
h.handleSortRoutes(w, r)
78+
return
79+
}
80+
6181
// GET /antigravity/providers/quotas - 批量获取配额（必须在单个 provider 路由之前匹配）
6282
if len(parts) >= 3 && parts[1] == "providers" && parts[2] == "quotas" && r.Method == http.MethodGet {
6383
h.handleGetBatchQuotas(w, r)
@@ -368,6 +388,8 @@ type BatchQuotaResult struct {
368388
}
369389

370390
// GetBatchQuotas 批量获取所有 Antigravity provider 的配额信息（供 HTTP handler 和 Wails 共用）
391+
// 优先从数据库返回缓存数据，即使过期也会返回（避免 API 请求阻塞）
392+
// 配额刷新由后台任务负责
371393
func (h *AntigravityHandler) GetBatchQuotas(ctx context.Context) (*BatchQuotaResult, error) {
372394
// 获取所有 providers
373395
providers, err := h.svc.GetProviders()
@@ -388,30 +410,19 @@ func (h *AntigravityHandler) GetBatchQuotas(ctx context.Context) (*BatchQuotaRes
388410
config := provider.Config.Antigravity
389411
email := config.Email
390412

391-
// 尝试从数据库获取缓存的配额
413+
// 优先从数据库获取缓存的配额（无论是否过期）
392414
if email != "" && h.quotaRepo != nil {
393415
cachedQuota, err := h.quotaRepo.GetByEmail(email)
394416
if err == nil && cachedQuota != nil {
395-
// 检查是否过期（10分钟）- 如果未过期，直接使用缓存
396-
if time.Since(cachedQuota.UpdatedAt).Seconds() < 600 {
397-
result.Quotas[provider.ID] = h.domainQuotaToResponse(cachedQuota)
398-
continue
399-
}
417+
result.Quotas[provider.ID] = h.domainQuotaToResponse(cachedQuota)
418+
continue
400419
}
401420
}
402421

403-
// 缓存过期或不存在，从 API 获取最新配额
422+
// 数据库没有缓存，尝试从 API 获取
404423
quota, err := antigravity.FetchQuotaForProvider(ctx, config.RefreshToken, config.ProjectID)
405424
if err != nil {
406-
// 如果 API 失败，尝试使用过期的缓存数据
407-
if email != "" && h.quotaRepo != nil {
408-
cachedQuota, _ := h.quotaRepo.GetByEmail(email)
409-
if cachedQuota != nil {
410-
result.Quotas[provider.ID] = h.domainQuotaToResponse(cachedQuota)
411-
continue
412-
}
413-
}
414-
// 跳过此 provider，不中断整体查询
425+
// API 失败，跳过此 provider
415426
continue
416427
}
417428

@@ -442,6 +453,33 @@ func (h *AntigravityHandler) handleGetBatchQuotas(w http.ResponseWriter, r *http
442453
writeJSON(w, http.StatusOK, result)
443454
}
444455

456+
// handleForceRefreshQuotas 强制刷新所有 Antigravity 配额
457+
func (h *AntigravityHandler) handleForceRefreshQuotas(w http.ResponseWriter, r *http.Request) {
458+
if h.taskSvc == nil {
459+
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "task service not available"})
460+
return
461+
}
462+
463+
refreshed := h.taskSvc.ForceRefreshQuotas(r.Context())
464+
writeJSON(w, http.StatusOK, map[string]interface{}{
465+
"success": true,
466+
"refreshed": refreshed,
467+
})
468+
}
469+
470+
// handleSortRoutes 手动排序 Antigravity 路由
471+
func (h *AntigravityHandler) handleSortRoutes(w http.ResponseWriter, r *http.Request) {
472+
if h.taskSvc == nil {
473+
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "task service not available"})
474+
return
475+
}
476+
477+
h.taskSvc.SortRoutes(r.Context())
478+
writeJSON(w, http.StatusOK, map[string]interface{}{
479+
"success": true,
480+
})
481+
}
482+
445483
// ============================================================================
446484
// OAuth 授权处理函数
447485
// ============================================================================

‎internal/repository/interfaces.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,8 @@ type ProxyRequestRepository interface {
7878
MarkStaleAsFailed(currentInstanceID string) (int64, error)
7979
// DeleteOlderThan 删除指定时间之前的请求记录
8080
DeleteOlderThan(before time.Time) (int64, error)
81+
// HasRecentRequests 检查指定时间之后是否有请求记录
82+
HasRecentRequests(since time.Time) (bool, error)
8183
}
8284

8385
type ProxyUpstreamAttemptRepository interface {

‎internal/repository/sqlite/proxy_request.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,16 @@ func (r *ProxyRequestRepository) DeleteOlderThan(before time.Time) (int64, error
187187
return affected, nil
188188
}
189189

190+
// HasRecentRequests 检查指定时间之后是否有请求记录
191+
func (r *ProxyRequestRepository) HasRecentRequests(since time.Time) (bool, error) {
192+
sinceTs := toTimestamp(since)
193+
var count int64
194+
if err := r.db.gorm.Model(&ProxyRequest{}).Where("created_at >= ?", sinceTs).Limit(1).Count(&count).Error; err != nil {
195+
return false, err
196+
}
197+
return count > 0, nil
198+
}
199+
190200
func (r *ProxyRequestRepository) toModel(p *domain.ProxyRequest) *ProxyRequest {
191201
return &ProxyRequest{
192202
BaseModel: BaseModel{

0 commit comments

Comments
 (0)