token_refresher.go 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. package service
  2. import (
  3. "context"
  4. "strings"
  5. "time"
  6. )
  7. // TokenRefresher 定义平台特定的token刷新策略接口
  8. // 通过此接口可以扩展支持不同平台(Anthropic/OpenAI/Gemini)
  9. type TokenRefresher interface {
  10. // CanRefresh 检查此刷新器是否能处理指定账号
  11. CanRefresh(account *Account) bool
  12. // NeedsRefresh 检查账号的token是否需要刷新
  13. NeedsRefresh(account *Account, refreshWindow time.Duration) bool
  14. // Refresh 执行token刷新,返回更新后的credentials
  15. // 注意:返回的map应该保留原有credentials中的所有字段,只更新token相关字段
  16. Refresh(ctx context.Context, account *Account) (map[string]any, error)
  17. }
  18. // ClaudeTokenRefresher 处理Anthropic/Claude OAuth token刷新
  19. type ClaudeTokenRefresher struct {
  20. oauthService *OAuthService
  21. }
  22. // NewClaudeTokenRefresher 创建Claude token刷新器
  23. func NewClaudeTokenRefresher(oauthService *OAuthService) *ClaudeTokenRefresher {
  24. return &ClaudeTokenRefresher{
  25. oauthService: oauthService,
  26. }
  27. }
  28. // CacheKey 返回用于分布式锁的缓存键
  29. func (r *ClaudeTokenRefresher) CacheKey(account *Account) string {
  30. return ClaudeTokenCacheKey(account)
  31. }
  32. // CanRefresh 检查是否能处理此账号
  33. // 处理 anthropic 平台的 oauth 与 setup-token 类型账号。
  34. // 两者的 access_token 均为短期令牌(expires_in=28800,即 8h),到期都需刷新;
  35. // setup-token 之前被排除会导致其 access_token 过期后请求 401。
  36. // 此处与手动刷新入口(account.IsOAuth())保持一致,实际是否刷新由 NeedsRefresh
  37. // 基于 expires_at 门控,并在分布式锁保护下执行,不会造成过度刷新。
  38. func (r *ClaudeTokenRefresher) CanRefresh(account *Account) bool {
  39. return account.Platform == PlatformAnthropic && account.IsOAuth()
  40. }
  41. // NeedsRefresh 检查token是否需要刷新
  42. // 基于 expires_at 字段判断是否在刷新窗口内
  43. func (r *ClaudeTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.Duration) bool {
  44. expiresAt := account.GetCredentialAsTime("expires_at")
  45. if expiresAt == nil {
  46. return false
  47. }
  48. return time.Until(*expiresAt) < refreshWindow
  49. }
  50. // Refresh 执行token刷新
  51. // 保留原有credentials中的所有字段,只更新token相关字段
  52. func (r *ClaudeTokenRefresher) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
  53. tokenInfo, err := r.oauthService.RefreshAccountToken(ctx, account)
  54. if err != nil {
  55. return nil, err
  56. }
  57. newCredentials := BuildClaudeAccountCredentials(tokenInfo)
  58. newCredentials = MergeCredentials(account.Credentials, newCredentials)
  59. return newCredentials, nil
  60. }
  61. // OpenAITokenRefresher 处理 OpenAI OAuth token刷新
  62. type OpenAITokenRefresher struct {
  63. openaiOAuthService *OpenAIOAuthService
  64. accountRepo AccountRepository
  65. }
  66. // NewOpenAITokenRefresher 创建 OpenAI token刷新器
  67. func NewOpenAITokenRefresher(openaiOAuthService *OpenAIOAuthService, accountRepo AccountRepository) *OpenAITokenRefresher {
  68. return &OpenAITokenRefresher{
  69. openaiOAuthService: openaiOAuthService,
  70. accountRepo: accountRepo,
  71. }
  72. }
  73. // CacheKey 返回用于分布式锁的缓存键
  74. func (r *OpenAITokenRefresher) CacheKey(account *Account) string {
  75. return OpenAITokenCacheKey(account)
  76. }
  77. // CanRefresh 检查是否能处理此账号
  78. func (r *OpenAITokenRefresher) CanRefresh(account *Account) bool {
  79. if account.IsCredentialShadow() {
  80. return false
  81. }
  82. return account.Platform == PlatformOpenAI && account.Type == AccountTypeOAuth
  83. }
  84. // NeedsRefresh 检查token是否需要刷新
  85. // expires_at 缺失且处于限流状态时需要刷新,防止限流期间 token 静默过期
  86. func (r *OpenAITokenRefresher) NeedsRefresh(account *Account, refreshWindow time.Duration) bool {
  87. if account.IsOpenAIPersonalAccessToken() {
  88. return false
  89. }
  90. if strings.TrimSpace(account.GetOpenAIRefreshToken()) == "" {
  91. return false
  92. }
  93. expiresAt := account.GetCredentialAsTime("expires_at")
  94. if expiresAt == nil {
  95. return account.IsRateLimited()
  96. }
  97. return time.Until(*expiresAt) < refreshWindow
  98. }
  99. // Refresh 执行token刷新
  100. // 保留原有credentials中的所有字段,只更新token相关字段
  101. func (r *OpenAITokenRefresher) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
  102. tokenInfo, err := r.openaiOAuthService.RefreshAccountToken(ctx, account)
  103. if err != nil {
  104. return nil, err
  105. }
  106. // 使用服务提供的方法构建新凭证,并保留原有字段
  107. newCredentials := r.openaiOAuthService.BuildAccountCredentials(tokenInfo)
  108. newCredentials = MergeCredentials(account.Credentials, newCredentials)
  109. newCredentials = NormalizeOpenAIPersonalAccessTokenCredentials(account, tokenInfo, newCredentials)
  110. return newCredentials, nil
  111. }