package service import ( "context" "log" "strconv" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" ) // AI Credits(GOOGLE_ONE_AI)状态管理: // - free tier 耗尽时,可注入 v1internal.enabledCreditTypes=["GOOGLE_ONE_AI"] 落到付费余额。 // - 账号余额持久化在 Account.Extra,避免每次请求都去 loadCodeAssist 查询。 // - INSUFFICIENT_G1_CREDITS_BALANCE 错误回写 exhausted_at 时间戳,让请求转换器跳过该账号的 credits 注入。 // - loadCodeAssist 时主动刷新余额并清除 exhausted 标记。 const ( // extraKeyCreditsBalance 缓存的 GOOGLE_ONE_AI 可用余额(float64,单位由上游决定)。 extraKeyCreditsBalance = "antigravity_credits_balance" // extraKeyCreditsCheckedAt 余额最近查询时间(RFC3339)。 extraKeyCreditsCheckedAt = "antigravity_credits_checked_at" // extraKeyCreditsExhaustedAt 上次收到 INSUFFICIENT_G1_CREDITS_BALANCE 的时间(RFC3339)。 extraKeyCreditsExhaustedAt = "antigravity_credits_exhausted_at" // creditsExhaustedRecheckInterval 余额耗尽后的重新探测间隔。 // 在此间隔内不再注入 enabledCreditTypes,避免反复触发 INSUFFICIENT_G1_CREDITS_BALANCE。 // 间隔到达后允许下一次 loadCodeAssist 刷新余额并解除标记。 creditsExhaustedRecheckInterval = 30 * time.Minute ) // AccountHasUsableCredits 判断账号当前是否可注入 enabledCreditTypes。 // - 仅当最近一次余额查询 > 0 且 exhausted_at 已过期才返回 true。 // - 从未查询过余额时返回 false(保守策略:不知道有就不注入,避免无效请求)。 func AccountHasUsableCredits(account *Account) bool { if account == nil || account.Extra == nil { return false } // 余额耗尽标记仍在生效期内 → 不可用 if exhaustedAtStr, ok := account.Extra[extraKeyCreditsExhaustedAt].(string); ok && exhaustedAtStr != "" { if t, err := time.Parse(time.RFC3339, exhaustedAtStr); err == nil { if time.Since(t) < creditsExhaustedRecheckInterval { return false } } } balance := readFloat(account.Extra[extraKeyCreditsBalance]) return balance > 0 } // refreshAccountCreditsFromLoadCodeAssist 从 loadCodeAssist 响应里提取 paidTier.availableCredits。 // 任何 GOOGLE_ONE_AI 类型的余额(或 creditType 为空时按总额)都会被写入 Account.Extra。 // 副作用:清除 exhausted 标记(因为我们刚刚拿到了上游确认的余额)。 func refreshAccountCreditsFromLoadCodeAssist(account *Account, resp *antigravity.LoadCodeAssistResponse) { if account == nil || resp == nil { return } if account.Extra == nil { account.Extra = make(map[string]any) } balance := pickGoogleOneAIBalance(resp.GetAvailableCredits()) account.Extra[extraKeyCreditsBalance] = balance account.Extra[extraKeyCreditsCheckedAt] = time.Now().UTC().Format(time.RFC3339) // 余额已重新查询;如果有余额或耗尽标记早于本次查询,应清除。 delete(account.Extra, extraKeyCreditsExhaustedAt) } // pickGoogleOneAIBalance 从 availableCredits 列表中提取 GOOGLE_ONE_AI 余额。 // creditType 为空(旧响应格式)时按整体余额累加。 func pickGoogleOneAIBalance(credits []antigravity.AvailableCredit) float64 { var total float64 for _, c := range credits { if c.CreditType == "" || c.CreditType == antigravity.CreditTypeGoogleOneAI { total += c.GetAmount() } } return total } // markAccountCreditsExhausted 把账号 credits 标记为余额不足(INSUFFICIENT_G1_CREDITS_BALANCE)。 // 余额改写为 0,写入耗尽时间戳,并把更新同步到 Redis 调度快照。 func (s *AntigravityGatewayService) markAccountCreditsExhausted(ctx context.Context, prefix string, account *Account) { if account == nil { return } if account.Extra == nil { account.Extra = make(map[string]any) } account.Extra[extraKeyCreditsBalance] = 0.0 account.Extra[extraKeyCreditsExhaustedAt] = time.Now().UTC().Format(time.RFC3339) account.Extra[extraKeyCreditsCheckedAt] = time.Now().UTC().Format(time.RFC3339) if s.schedulerSnapshot != nil { if err := s.schedulerSnapshot.UpdateAccountInCache(ctx, account); err != nil { log.Printf("%s credits_exhausted_cache_update_failed account=%d err=%v", prefix, account.ID, err) } } } // readFloat 从 Account.Extra 的 any 类型里宽松读取浮点数(兼容 JSON 反序列化的 float64/string)。 func readFloat(v any) float64 { switch x := v.(type) { case float64: return x case float32: return float64(x) case int: return float64(x) case int64: return float64(x) case string: if x == "" { return 0 } f, err := strconv.ParseFloat(x, 64) if err != nil { return 0 } return f } return 0 }