antigravity_token_refresher.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. package service
  2. import (
  3. "context"
  4. "fmt"
  5. "log"
  6. "strings"
  7. "time"
  8. )
  9. const (
  10. // antigravityRefreshWindow Antigravity token 提前刷新窗口:15分钟
  11. // Google OAuth token 有效期55分钟,提前15分钟刷新
  12. antigravityRefreshWindow = 15 * time.Minute
  13. antigravityForceTokenRefreshExtraKey = "antigravity_force_token_refresh"
  14. antigravityForceTokenRefreshReasonExtraKey = "antigravity_force_token_refresh_reason"
  15. antigravityForceTokenRefreshAtExtraKey = "antigravity_force_token_refresh_at"
  16. )
  17. // AntigravityTokenRefresher 实现 TokenRefresher 接口
  18. type AntigravityTokenRefresher struct {
  19. antigravityOAuthService *AntigravityOAuthService
  20. }
  21. func NewAntigravityTokenRefresher(antigravityOAuthService *AntigravityOAuthService) *AntigravityTokenRefresher {
  22. return &AntigravityTokenRefresher{
  23. antigravityOAuthService: antigravityOAuthService,
  24. }
  25. }
  26. // CacheKey 返回用于分布式锁的缓存键
  27. func (r *AntigravityTokenRefresher) CacheKey(account *Account) string {
  28. return AntigravityTokenCacheKey(account)
  29. }
  30. // CanRefresh 检查是否可以刷新此账户
  31. func (r *AntigravityTokenRefresher) CanRefresh(account *Account) bool {
  32. return account.Platform == PlatformAntigravity && account.Type == AccountTypeOAuth
  33. }
  34. // NeedsRefresh 检查账户是否需要刷新
  35. // Antigravity 使用固定的15分钟刷新窗口,忽略全局配置
  36. func (r *AntigravityTokenRefresher) NeedsRefresh(account *Account, _ time.Duration) bool {
  37. if !r.CanRefresh(account) {
  38. return false
  39. }
  40. if accountNeedsAntigravityForceTokenRefresh(account) {
  41. return true
  42. }
  43. expiresAt := account.GetCredentialAsTime("expires_at")
  44. if expiresAt == nil {
  45. return false
  46. }
  47. timeUntilExpiry := time.Until(*expiresAt)
  48. needsRefresh := timeUntilExpiry < antigravityRefreshWindow
  49. if needsRefresh {
  50. fmt.Printf("[AntigravityTokenRefresher] Account %d needs refresh: expires_at=%s, time_until_expiry=%v, window=%v\n",
  51. account.ID, expiresAt.Format("2006-01-02 15:04:05"), timeUntilExpiry, antigravityRefreshWindow)
  52. }
  53. return needsRefresh
  54. }
  55. func accountNeedsAntigravityForceTokenRefresh(account *Account) bool {
  56. return account != nil &&
  57. account.Platform == PlatformAntigravity &&
  58. account.Type == AccountTypeOAuth &&
  59. account.getExtraBool(antigravityForceTokenRefreshExtraKey)
  60. }
  61. func antigravityForceTokenRefreshExtra(reason string) map[string]any {
  62. return map[string]any{
  63. antigravityForceTokenRefreshExtraKey: true,
  64. antigravityForceTokenRefreshReasonExtraKey: reason,
  65. antigravityForceTokenRefreshAtExtraKey: time.Now().UTC().Format(time.RFC3339),
  66. }
  67. }
  68. func clearAntigravityForceTokenRefreshExtra() map[string]any {
  69. return map[string]any{
  70. antigravityForceTokenRefreshExtraKey: false,
  71. antigravityForceTokenRefreshReasonExtraKey: "",
  72. antigravityForceTokenRefreshAtExtraKey: "",
  73. }
  74. }
  75. // Refresh 执行 token 刷新
  76. func (r *AntigravityTokenRefresher) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
  77. tokenInfo, err := r.antigravityOAuthService.RefreshAccountToken(ctx, account)
  78. if err != nil {
  79. return nil, err
  80. }
  81. newCredentials := r.antigravityOAuthService.BuildAccountCredentials(tokenInfo)
  82. // 合并旧的 credentials,保留新 credentials 中不存在的字段
  83. newCredentials = MergeCredentials(account.Credentials, newCredentials)
  84. // 特殊处理 project_id:如果新值为空但旧值非空,保留旧值
  85. // 这确保了即使 LoadCodeAssist 失败,project_id 也不会丢失
  86. if newProjectID, _ := newCredentials["project_id"].(string); newProjectID == "" {
  87. if oldProjectID := strings.TrimSpace(account.GetCredential("project_id")); oldProjectID != "" {
  88. newCredentials["project_id"] = oldProjectID
  89. }
  90. }
  91. // 如果 project_id 获取失败,只记录警告,不返回错误
  92. // LoadCodeAssist 失败可能是临时网络问题,应该允许重试而不是立即标记为不可重试错误
  93. // Token 刷新本身是成功的(access_token 和 refresh_token 已更新)
  94. if tokenInfo.ProjectIDMissing {
  95. if tokenInfo.ProjectID != "" {
  96. // 有旧的 project_id,本次获取失败,保留旧值
  97. log.Printf("[AntigravityTokenRefresher] Account %d: LoadCodeAssist 临时失败,保留旧 project_id", account.ID)
  98. } else {
  99. // 从未获取过 project_id,本次也失败,但不返回错误以允许下次重试
  100. log.Printf("[AntigravityTokenRefresher] Account %d: LoadCodeAssist 失败,project_id 缺失,但 token 已更新,将在下次刷新时重试", account.ID)
  101. }
  102. }
  103. return newCredentials, nil
  104. }