You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

83 lines
2.1 KiB

  1. package service
  2. import (
  3. "context"
  4. "fmt"
  5. "strings"
  6. "time"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/constant"
  9. "github.com/QuantumNous/new-api/model"
  10. )
  11. type CodexCredentialRefreshOptions struct {
  12. ResetCaches bool
  13. }
  14. type CodexOAuthKey = common.CodexOAuthCredential
  15. func RefreshCodexChannelCredential(ctx context.Context, channelID int, opts CodexCredentialRefreshOptions) (*CodexOAuthKey, *model.Channel, error) {
  16. ch, err := model.GetChannelById(channelID, true)
  17. if err != nil {
  18. return nil, nil, err
  19. }
  20. if ch == nil {
  21. return nil, nil, fmt.Errorf("channel not found")
  22. }
  23. if ch.Type != constant.ChannelTypeCodex {
  24. return nil, nil, fmt.Errorf("channel type is not Codex")
  25. }
  26. oauthKey, err := common.ParseCodexOAuthCredential(strings.TrimSpace(ch.Key))
  27. if err != nil {
  28. return nil, nil, err
  29. }
  30. if strings.TrimSpace(oauthKey.RefreshToken) == "" {
  31. return nil, nil, fmt.Errorf("codex channel: refresh_token is required to refresh credential")
  32. }
  33. refreshCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
  34. defer cancel()
  35. res, err := RefreshCodexOAuthTokenWithProxy(refreshCtx, oauthKey.RefreshToken, ch.GetSetting().Proxy)
  36. if err != nil {
  37. return nil, nil, err
  38. }
  39. oauthKey.AccessToken = res.AccessToken
  40. oauthKey.RefreshToken = res.RefreshToken
  41. oauthKey.LastRefresh = time.Now().Format(time.RFC3339)
  42. oauthKey.Expired = res.ExpiresAt.Format(time.RFC3339)
  43. if strings.TrimSpace(oauthKey.Type) == "" {
  44. oauthKey.Type = "codex"
  45. }
  46. if strings.TrimSpace(oauthKey.AccountID) == "" {
  47. if accountID, ok := ExtractCodexAccountIDFromJWT(oauthKey.AccessToken); ok {
  48. oauthKey.AccountID = accountID
  49. }
  50. }
  51. if strings.TrimSpace(oauthKey.Email) == "" {
  52. if email, ok := ExtractEmailFromJWT(oauthKey.AccessToken); ok {
  53. oauthKey.Email = email
  54. }
  55. }
  56. encoded, err := common.Marshal(oauthKey)
  57. if err != nil {
  58. return nil, nil, err
  59. }
  60. if err := model.DB.Model(&model.Channel{}).Where("id = ?", ch.Id).Update("key", string(encoded)).Error; err != nil {
  61. return nil, nil, err
  62. }
  63. if opts.ResetCaches {
  64. model.InitChannelCache()
  65. ResetProxyClientCache()
  66. }
  67. return oauthKey, ch, nil
  68. }