token_cache_invalidator.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  1. package service
  2. import (
  3. "context"
  4. "log/slog"
  5. "strconv"
  6. )
  7. type TokenCacheInvalidator interface {
  8. InvalidateToken(ctx context.Context, account *Account) error
  9. }
  10. type CompositeTokenCacheInvalidator struct {
  11. cache GeminiTokenCache // 统一使用一个缓存接口,通过缓存键前缀区分平台
  12. }
  13. func NewCompositeTokenCacheInvalidator(cache GeminiTokenCache) *CompositeTokenCacheInvalidator {
  14. return &CompositeTokenCacheInvalidator{
  15. cache: cache,
  16. }
  17. }
  18. func (c *CompositeTokenCacheInvalidator) InvalidateToken(ctx context.Context, account *Account) error {
  19. if c == nil || c.cache == nil || account == nil {
  20. return nil
  21. }
  22. if account.Type != AccountTypeOAuth {
  23. return nil
  24. }
  25. var keysToDelete []string
  26. accountIDKey := "account:" + strconv.FormatInt(account.ID, 10)
  27. switch account.Platform {
  28. case PlatformGemini:
  29. // Gemini 可能有两种缓存键:project_id 或 account_id
  30. // 首次获取 token 时可能没有 project_id,之后自动检测到 project_id 后会使用新 key
  31. // 刷新时需要同时删除两种可能的 key,确保不会遗留旧缓存
  32. keysToDelete = append(keysToDelete, GeminiTokenCacheKey(account))
  33. keysToDelete = append(keysToDelete, "gemini:"+accountIDKey)
  34. case PlatformAntigravity:
  35. // Antigravity 同样可能有两种缓存键
  36. keysToDelete = append(keysToDelete, AntigravityTokenCacheKey(account))
  37. keysToDelete = append(keysToDelete, "ag:"+accountIDKey)
  38. case PlatformOpenAI:
  39. keysToDelete = append(keysToDelete, OpenAITokenCacheKey(account))
  40. case PlatformGrok:
  41. keysToDelete = append(keysToDelete, GrokTokenCacheKey(account))
  42. keysToDelete = append(keysToDelete, "grok:"+accountIDKey)
  43. case PlatformAnthropic:
  44. keysToDelete = append(keysToDelete, ClaudeTokenCacheKey(account))
  45. default:
  46. return nil
  47. }
  48. // 删除所有可能的缓存键(去重后)
  49. seen := make(map[string]bool)
  50. for _, key := range keysToDelete {
  51. if seen[key] {
  52. continue
  53. }
  54. seen[key] = true
  55. if err := c.cache.DeleteAccessToken(ctx, key); err != nil {
  56. slog.Warn("token_cache_delete_failed", "key", key, "account_id", account.ID, "error", err)
  57. }
  58. }
  59. return nil
  60. }
  61. // CheckTokenVersion 检查 account 的 token 版本是否已过时,并返回最新的 account
  62. // 用于解决异步刷新任务与请求线程的竞态条件:
  63. // 如果刷新任务已更新 token 并删除缓存,此时请求线程的旧 account 对象不应写入缓存
  64. //
  65. // 返回值:
  66. // - latestAccount: 从 DB 获取的最新 account(如果查询失败则返回 nil)
  67. // - isStale: true 表示 token 已过时(应使用 latestAccount),false 表示可以使用当前 account
  68. func CheckTokenVersion(ctx context.Context, account *Account, repo AccountRepository) (latestAccount *Account, isStale bool) {
  69. if account == nil || repo == nil {
  70. return nil, false
  71. }
  72. currentVersion := account.GetCredentialAsInt64("_token_version")
  73. latestAccount, err := repo.GetByID(ctx, account.ID)
  74. if err != nil || latestAccount == nil {
  75. // 查询失败,默认允许缓存,不返回 latestAccount
  76. return nil, false
  77. }
  78. latestVersion := latestAccount.GetCredentialAsInt64("_token_version")
  79. // 情况1: 当前 account 没有版本号,但 DB 中已有版本号
  80. // 说明异步刷新任务已更新 token,当前 account 已过时
  81. if currentVersion == 0 && latestVersion > 0 {
  82. slog.Debug("token_version_stale_no_current_version",
  83. "account_id", account.ID,
  84. "latest_version", latestVersion)
  85. return latestAccount, true
  86. }
  87. // 情况2: 两边都没有版本号,说明从未被异步刷新过,允许缓存
  88. if currentVersion == 0 && latestVersion == 0 {
  89. return latestAccount, false
  90. }
  91. // 情况3: 比较版本号,如果 DB 中的版本更新,当前 account 已过时
  92. if latestVersion > currentVersion {
  93. slog.Debug("token_version_stale",
  94. "account_id", account.ID,
  95. "current_version", currentVersion,
  96. "latest_version", latestVersion)
  97. return latestAccount, true
  98. }
  99. return latestAccount, false
  100. }