新增 Account.GetCredentialAsTime 方法,统一处理凭证中的时间戳字段, 兼容 RFC3339 字符串、Unix 时间戳字符串和数字类型。 - 重构 Claude/Gemini/Antigravity TokenRefresher.NeedsRefresh - 移除重复的 parseExpiresAt/parseAntigravityExpiresAt 函数 - 简化 GetOpenAITokenExpiresAt 实现 - 新增 RFC3339 格式单元测试用例
376 lines
7.5 KiB
Go
376 lines
7.5 KiB
Go
package service
|
|
|
|
import (
|
|
"encoding/json"
|
|
"strconv"
|
|
"time"
|
|
)
|
|
|
|
type Account struct {
|
|
ID int64
|
|
Name string
|
|
Platform string
|
|
Type string
|
|
Credentials map[string]any
|
|
Extra map[string]any
|
|
ProxyID *int64
|
|
Concurrency int
|
|
Priority int
|
|
Status string
|
|
ErrorMessage string
|
|
LastUsedAt *time.Time
|
|
CreatedAt time.Time
|
|
UpdatedAt time.Time
|
|
|
|
Schedulable bool
|
|
|
|
RateLimitedAt *time.Time
|
|
RateLimitResetAt *time.Time
|
|
OverloadUntil *time.Time
|
|
|
|
SessionWindowStart *time.Time
|
|
SessionWindowEnd *time.Time
|
|
SessionWindowStatus string
|
|
|
|
Proxy *Proxy
|
|
AccountGroups []AccountGroup
|
|
GroupIDs []int64
|
|
Groups []*Group
|
|
}
|
|
|
|
func (a *Account) IsActive() bool {
|
|
return a.Status == StatusActive
|
|
}
|
|
|
|
func (a *Account) IsSchedulable() bool {
|
|
if !a.IsActive() || !a.Schedulable {
|
|
return false
|
|
}
|
|
now := time.Now()
|
|
if a.OverloadUntil != nil && now.Before(*a.OverloadUntil) {
|
|
return false
|
|
}
|
|
if a.RateLimitResetAt != nil && now.Before(*a.RateLimitResetAt) {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (a *Account) IsRateLimited() bool {
|
|
if a.RateLimitResetAt == nil {
|
|
return false
|
|
}
|
|
return time.Now().Before(*a.RateLimitResetAt)
|
|
}
|
|
|
|
func (a *Account) IsOverloaded() bool {
|
|
if a.OverloadUntil == nil {
|
|
return false
|
|
}
|
|
return time.Now().Before(*a.OverloadUntil)
|
|
}
|
|
|
|
func (a *Account) IsOAuth() bool {
|
|
return a.Type == AccountTypeOAuth || a.Type == AccountTypeSetupToken
|
|
}
|
|
|
|
func (a *Account) IsGemini() bool {
|
|
return a.Platform == PlatformGemini
|
|
}
|
|
|
|
func (a *Account) CanGetUsage() bool {
|
|
return a.Type == AccountTypeOAuth
|
|
}
|
|
|
|
func (a *Account) GetCredential(key string) string {
|
|
if a.Credentials == nil {
|
|
return ""
|
|
}
|
|
v, ok := a.Credentials[key]
|
|
if !ok || v == nil {
|
|
return ""
|
|
}
|
|
|
|
// 支持多种类型(兼容历史数据中 expires_at 等字段可能是数字或字符串)
|
|
switch val := v.(type) {
|
|
case string:
|
|
return val
|
|
case json.Number:
|
|
// GORM datatypes.JSONMap 使用 UseNumber() 解析,数字类型为 json.Number
|
|
return val.String()
|
|
case float64:
|
|
// JSON 解析后数字默认为 float64
|
|
return strconv.FormatInt(int64(val), 10)
|
|
case int64:
|
|
return strconv.FormatInt(val, 10)
|
|
case int:
|
|
return strconv.Itoa(val)
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
// GetCredentialAsTime 解析凭证中的时间戳字段,支持多种格式
|
|
// 兼容以下格式:
|
|
// - RFC3339 字符串: "2025-01-01T00:00:00Z"
|
|
// - Unix 时间戳字符串: "1735689600"
|
|
// - Unix 时间戳数字: 1735689600 (float64/int64/json.Number)
|
|
func (a *Account) GetCredentialAsTime(key string) *time.Time {
|
|
s := a.GetCredential(key)
|
|
if s == "" {
|
|
return nil
|
|
}
|
|
// 尝试 RFC3339 格式
|
|
if t, err := time.Parse(time.RFC3339, s); err == nil {
|
|
return &t
|
|
}
|
|
// 尝试 Unix 时间戳(纯数字字符串)
|
|
if ts, err := strconv.ParseInt(s, 10, 64); err == nil {
|
|
t := time.Unix(ts, 0)
|
|
return &t
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (a *Account) GetModelMapping() map[string]string {
|
|
if a.Credentials == nil {
|
|
return nil
|
|
}
|
|
raw, ok := a.Credentials["model_mapping"]
|
|
if !ok || raw == nil {
|
|
return nil
|
|
}
|
|
if m, ok := raw.(map[string]any); ok {
|
|
result := make(map[string]string)
|
|
for k, v := range m {
|
|
if s, ok := v.(string); ok {
|
|
result[k] = s
|
|
}
|
|
}
|
|
if len(result) > 0 {
|
|
return result
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (a *Account) IsModelSupported(requestedModel string) bool {
|
|
mapping := a.GetModelMapping()
|
|
if len(mapping) == 0 {
|
|
return true
|
|
}
|
|
_, exists := mapping[requestedModel]
|
|
return exists
|
|
}
|
|
|
|
func (a *Account) GetMappedModel(requestedModel string) string {
|
|
mapping := a.GetModelMapping()
|
|
if len(mapping) == 0 {
|
|
return requestedModel
|
|
}
|
|
if mappedModel, exists := mapping[requestedModel]; exists {
|
|
return mappedModel
|
|
}
|
|
return requestedModel
|
|
}
|
|
|
|
func (a *Account) GetBaseURL() string {
|
|
if a.Type != AccountTypeApiKey {
|
|
return ""
|
|
}
|
|
baseURL := a.GetCredential("base_url")
|
|
if baseURL == "" {
|
|
return "https://api.anthropic.com"
|
|
}
|
|
return baseURL
|
|
}
|
|
|
|
func (a *Account) GetExtraString(key string) string {
|
|
if a.Extra == nil {
|
|
return ""
|
|
}
|
|
if v, ok := a.Extra[key]; ok {
|
|
if s, ok := v.(string); ok {
|
|
return s
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (a *Account) IsCustomErrorCodesEnabled() bool {
|
|
if a.Type != AccountTypeApiKey || a.Credentials == nil {
|
|
return false
|
|
}
|
|
if v, ok := a.Credentials["custom_error_codes_enabled"]; ok {
|
|
if enabled, ok := v.(bool); ok {
|
|
return enabled
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (a *Account) GetCustomErrorCodes() []int {
|
|
if a.Credentials == nil {
|
|
return nil
|
|
}
|
|
raw, ok := a.Credentials["custom_error_codes"]
|
|
if !ok || raw == nil {
|
|
return nil
|
|
}
|
|
if arr, ok := raw.([]any); ok {
|
|
result := make([]int, 0, len(arr))
|
|
for _, v := range arr {
|
|
if f, ok := v.(float64); ok {
|
|
result = append(result, int(f))
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (a *Account) ShouldHandleErrorCode(statusCode int) bool {
|
|
if !a.IsCustomErrorCodesEnabled() {
|
|
return true
|
|
}
|
|
codes := a.GetCustomErrorCodes()
|
|
if len(codes) == 0 {
|
|
return true
|
|
}
|
|
for _, code := range codes {
|
|
if code == statusCode {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (a *Account) IsInterceptWarmupEnabled() bool {
|
|
if a.Credentials == nil {
|
|
return false
|
|
}
|
|
if v, ok := a.Credentials["intercept_warmup_requests"]; ok {
|
|
if enabled, ok := v.(bool); ok {
|
|
return enabled
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (a *Account) IsOpenAI() bool {
|
|
return a.Platform == PlatformOpenAI
|
|
}
|
|
|
|
func (a *Account) IsAnthropic() bool {
|
|
return a.Platform == PlatformAnthropic
|
|
}
|
|
|
|
func (a *Account) IsOpenAIOAuth() bool {
|
|
return a.IsOpenAI() && a.Type == AccountTypeOAuth
|
|
}
|
|
|
|
func (a *Account) IsOpenAIApiKey() bool {
|
|
return a.IsOpenAI() && a.Type == AccountTypeApiKey
|
|
}
|
|
|
|
func (a *Account) GetOpenAIBaseURL() string {
|
|
if !a.IsOpenAI() {
|
|
return ""
|
|
}
|
|
if a.Type == AccountTypeApiKey {
|
|
baseURL := a.GetCredential("base_url")
|
|
if baseURL != "" {
|
|
return baseURL
|
|
}
|
|
}
|
|
return "https://api.openai.com"
|
|
}
|
|
|
|
func (a *Account) GetOpenAIAccessToken() string {
|
|
if !a.IsOpenAI() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("access_token")
|
|
}
|
|
|
|
func (a *Account) GetOpenAIRefreshToken() string {
|
|
if !a.IsOpenAIOAuth() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("refresh_token")
|
|
}
|
|
|
|
func (a *Account) GetOpenAIIDToken() string {
|
|
if !a.IsOpenAIOAuth() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("id_token")
|
|
}
|
|
|
|
func (a *Account) GetOpenAIApiKey() string {
|
|
if !a.IsOpenAIApiKey() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("api_key")
|
|
}
|
|
|
|
func (a *Account) GetOpenAIUserAgent() string {
|
|
if !a.IsOpenAI() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("user_agent")
|
|
}
|
|
|
|
func (a *Account) GetChatGPTAccountID() string {
|
|
if !a.IsOpenAIOAuth() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("chatgpt_account_id")
|
|
}
|
|
|
|
func (a *Account) GetChatGPTUserID() string {
|
|
if !a.IsOpenAIOAuth() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("chatgpt_user_id")
|
|
}
|
|
|
|
func (a *Account) GetOpenAIOrganizationID() string {
|
|
if !a.IsOpenAIOAuth() {
|
|
return ""
|
|
}
|
|
return a.GetCredential("organization_id")
|
|
}
|
|
|
|
func (a *Account) GetOpenAITokenExpiresAt() *time.Time {
|
|
if !a.IsOpenAIOAuth() {
|
|
return nil
|
|
}
|
|
return a.GetCredentialAsTime("expires_at")
|
|
}
|
|
|
|
func (a *Account) IsOpenAITokenExpired() bool {
|
|
expiresAt := a.GetOpenAITokenExpiresAt()
|
|
if expiresAt == nil {
|
|
return false
|
|
}
|
|
return time.Now().Add(60 * time.Second).After(*expiresAt)
|
|
}
|
|
|
|
// IsMixedSchedulingEnabled 检查 antigravity 账户是否启用混合调度
|
|
// 启用后可参与 anthropic/gemini 分组的账户调度
|
|
func (a *Account) IsMixedSchedulingEnabled() bool {
|
|
if a.Platform != PlatformAntigravity {
|
|
return false
|
|
}
|
|
if a.Extra == nil {
|
|
return false
|
|
}
|
|
if v, ok := a.Extra["mixed_scheduling"]; ok {
|
|
if enabled, ok := v.(bool); ok {
|
|
return enabled
|
|
}
|
|
}
|
|
return false
|
|
}
|