feat(backend): add kiro account support
This commit is contained in:
nianzs
@@ -0,0 +1,608 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
|
||||
kiropkg "github.com/Wei-Shaw/sub2api/internal/pkg/kiro"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
kiroUsageOrigin = "AI_EDITOR"
|
||||
kiroUsageResourceType = "AGENTIC_REQUEST"
|
||||
|
||||
kiroDefaultRegion = "us-east-1"
|
||||
)
|
||||
|
||||
var resolveKiroRuntimeEndpoint = kiroRuntimeEndpoint
|
||||
|
||||
type kiroUsageLimitsResponse struct {
|
||||
NextDateReset any `json:"nextDateReset"`
|
||||
OverageConfiguration kiroOverageConfiguration `json:"overageConfiguration"`
|
||||
SubscriptionInfo kiroSubscriptionInfo `json:"subscriptionInfo"`
|
||||
UsageBreakdownList []kiroUsageBreakdown `json:"usageBreakdownList"`
|
||||
}
|
||||
|
||||
type kiroOverageConfiguration struct {
|
||||
OverageStatus string `json:"overageStatus"`
|
||||
}
|
||||
|
||||
type kiroSubscriptionInfo struct {
|
||||
SubscriptionTitle string `json:"subscriptionTitle"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
type kiroUsageBreakdown struct {
|
||||
Currency string `json:"currency"`
|
||||
CurrentOverages *float64 `json:"currentOverages"`
|
||||
CurrentOveragesWithPrecision *float64 `json:"currentOveragesWithPrecision"`
|
||||
CurrentUsage *float64 `json:"currentUsage"`
|
||||
CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision"`
|
||||
DisplayName string `json:"displayName"`
|
||||
DisplayNamePlural string `json:"displayNamePlural"`
|
||||
FreeTrialInfo *kiroFreeTrialInfo `json:"freeTrialInfo"`
|
||||
NextDateReset any `json:"nextDateReset"`
|
||||
OverageCharges *float64 `json:"overageCharges"`
|
||||
ResourceType string `json:"resourceType"`
|
||||
UsageLimit *float64 `json:"usageLimit"`
|
||||
UsageLimitWithPrecision *float64 `json:"usageLimitWithPrecision"`
|
||||
}
|
||||
|
||||
type kiroFreeTrialInfo struct {
|
||||
CurrentUsage *float64 `json:"currentUsage"`
|
||||
CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision"`
|
||||
FreeTrialExpiry any `json:"freeTrialExpiry"`
|
||||
FreeTrialStatus string `json:"freeTrialStatus"`
|
||||
UsageLimit *float64 `json:"usageLimit"`
|
||||
UsageLimitWithPrecision *float64 `json:"usageLimitWithPrecision"`
|
||||
}
|
||||
|
||||
type kiroUsageHTTPError struct {
|
||||
StatusCode int
|
||||
Body string
|
||||
}
|
||||
|
||||
func (e *kiroUsageHTTPError) Error() string {
|
||||
if e == nil {
|
||||
return "kiro usage request failed"
|
||||
}
|
||||
if strings.TrimSpace(e.Body) == "" {
|
||||
return fmt.Sprintf("kiro usage request failed (status %d)", e.StatusCode)
|
||||
}
|
||||
return fmt.Sprintf("kiro usage request failed (status %d): %s", e.StatusCode, e.Body)
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) getKiroUsage(ctx context.Context, account *Account, source string, forceRefresh bool) (*UsageInfo, error) {
|
||||
now := time.Now()
|
||||
if account == nil {
|
||||
return &UsageInfo{
|
||||
Source: source,
|
||||
UpdatedAt: &now,
|
||||
Error: "account is nil",
|
||||
ErrorCode: errorCodeNetworkError,
|
||||
}, nil
|
||||
}
|
||||
if account.Platform != PlatformKiro || account.Type != AccountTypeOAuth {
|
||||
return &UsageInfo{
|
||||
Source: source,
|
||||
UpdatedAt: &now,
|
||||
}, nil
|
||||
}
|
||||
|
||||
cached, hasCached := s.getCachedKiroUsage(account.ID)
|
||||
if hasCached && (cached.ErrorCode != "" || cached.Error != "") {
|
||||
cached.Source = source
|
||||
s.attachKiroRuntimeState(ctx, account, cached)
|
||||
return cached, nil
|
||||
}
|
||||
if !forceRefresh && hasCached {
|
||||
cached.Source = source
|
||||
s.attachKiroRuntimeState(ctx, account, cached)
|
||||
return cached, nil
|
||||
}
|
||||
|
||||
flightKey := fmt.Sprintf("kiro-usage:%d", account.ID)
|
||||
result, fetchErr, _ := s.cache.kiroUsageFlight.Do(flightKey, func() (any, error) {
|
||||
if !forceRefresh {
|
||||
if usage, ok := s.getCachedKiroUsage(account.ID); ok {
|
||||
return usage, nil
|
||||
}
|
||||
}
|
||||
usage, err := s.fetchAndCacheKiroUsage(ctx, account, source)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return usage, nil
|
||||
})
|
||||
if fetchErr == nil {
|
||||
if usage, ok := result.(*UsageInfo); ok && usage != nil {
|
||||
usage.Source = source
|
||||
s.attachKiroRuntimeState(ctx, account, usage)
|
||||
if source == "active" {
|
||||
s.tryClearRecoverableAccountError(ctx, account)
|
||||
}
|
||||
return usage, nil
|
||||
}
|
||||
}
|
||||
|
||||
degraded := buildKiroDegradedUsage(fetchErr)
|
||||
degraded.Source = source
|
||||
if hasCached {
|
||||
cached.Error = degraded.Error
|
||||
cached.ErrorCode = degraded.ErrorCode
|
||||
cached.NeedsReauth = degraded.NeedsReauth
|
||||
cached.KiroQuotaState = degraded.KiroQuotaState
|
||||
cached.KiroQuotaReason = degraded.KiroQuotaReason
|
||||
cached.KiroQuotaResetAt = degraded.KiroQuotaResetAt
|
||||
cached.Source = source
|
||||
s.attachKiroRuntimeState(ctx, account, cached)
|
||||
return cached, nil
|
||||
}
|
||||
s.storeKiroUsageSnapshot(account.ID, degraded)
|
||||
s.attachKiroRuntimeState(ctx, account, degraded)
|
||||
return degraded, nil
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) fetchAndCacheKiroUsage(ctx context.Context, account *Account, source string) (*UsageInfo, error) {
|
||||
token := strings.TrimSpace(account.GetCredential("access_token"))
|
||||
if token == "" {
|
||||
return nil, fmt.Errorf("no access token available")
|
||||
}
|
||||
|
||||
region := kiroAPIRegion(account)
|
||||
profileArn := strings.TrimSpace(account.GetCredential("profile_arn"))
|
||||
|
||||
resp, err := s.requestKiroUsageLimits(ctx, account, region, profileArn, token)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
usage := mapKiroUsageToInfo(resp)
|
||||
usage.Source = source
|
||||
s.storeKiroUsageSnapshot(account.ID, usage)
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) storeKiroUsageSnapshot(accountID int64, usage *UsageInfo) {
|
||||
if s == nil || s.cache == nil || accountID <= 0 || usage == nil {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
if usage.UpdatedAt == nil {
|
||||
usage.UpdatedAt = &now
|
||||
}
|
||||
s.cache.kiroUsageCache.Store(accountID, &kiroUsageCache{
|
||||
usageInfo: cloneUsageInfo(usage),
|
||||
timestamp: now,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) getCachedKiroUsage(accountID int64) (*UsageInfo, bool) {
|
||||
if s == nil || s.cache == nil || accountID <= 0 {
|
||||
return nil, false
|
||||
}
|
||||
cached, ok := s.cache.kiroUsageCache.Load(accountID)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
cache, ok := cached.(*kiroUsageCache)
|
||||
if !ok || cache == nil || cache.usageInfo == nil {
|
||||
return nil, false
|
||||
}
|
||||
if time.Since(cache.timestamp) >= kiroCacheTTL(cache.usageInfo) {
|
||||
return nil, false
|
||||
}
|
||||
return cloneUsageInfo(cache.usageInfo), true
|
||||
}
|
||||
|
||||
func kiroCacheTTL(info *UsageInfo) time.Duration {
|
||||
if info == nil {
|
||||
return kiroUsageErrorTTL
|
||||
}
|
||||
if info.ErrorCode != "" || info.Error != "" {
|
||||
return kiroUsageErrorTTL
|
||||
}
|
||||
return apiCacheTTL
|
||||
}
|
||||
|
||||
func cloneUsageInfo(info *UsageInfo) *UsageInfo {
|
||||
if info == nil {
|
||||
return nil
|
||||
}
|
||||
cloned := *info
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) requestKiroUsageLimits(ctx context.Context, account *Account, region, profileArn, token string) (*kiroUsageLimitsResponse, error) {
|
||||
endpoint := resolveKiroRuntimeEndpoint(region)
|
||||
reqURL, err := url.Parse(endpoint + "/getUsageLimits")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build kiro usage url failed: %w", err)
|
||||
}
|
||||
q := reqURL.Query()
|
||||
q.Set("origin", kiroUsageOrigin)
|
||||
if profileArn = strings.TrimSpace(profileArn); profileArn != "" {
|
||||
q.Set("profileArn", profileArn)
|
||||
}
|
||||
q.Set("resourceType", kiroUsageResourceType)
|
||||
reqURL.RawQuery = q.Encode()
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL.String(), nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create kiro usage request failed: %w", err)
|
||||
}
|
||||
s.applyKiroRuntimeHeaders(req, account, token)
|
||||
|
||||
client, err := httpclient.GetClient(httpclient.Options{
|
||||
ProxyURL: accountProxyURL(account),
|
||||
Timeout: 30 * time.Second,
|
||||
ValidateResolvedIP: true,
|
||||
AllowPrivateHosts: isLoopbackEndpoint(endpoint),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create kiro usage client failed: %w", err)
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("kiro usage request failed: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read kiro usage response failed: %w", err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, &kiroUsageHTTPError{StatusCode: resp.StatusCode, Body: strings.TrimSpace(string(body))}
|
||||
}
|
||||
|
||||
var parsed kiroUsageLimitsResponse
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("decode kiro usage response failed: %w", err)
|
||||
}
|
||||
return &parsed, nil
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) applyKiroRuntimeHeaders(req *http.Request, account *Account, token string) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
accountKey := buildKiroAccountKey(account)
|
||||
machineID := buildKiroMachineID(account)
|
||||
req.Header.Set("Accept", "*/*")
|
||||
req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(token))
|
||||
req.Header.Set("User-Agent", kiropkg.BuildRuntimeUserAgent(accountKey, machineID))
|
||||
req.Header.Set("X-Amz-User-Agent", kiropkg.BuildRuntimeAmzUserAgent(accountKey, machineID))
|
||||
req.Header.Set("x-amzn-kiro-agent-mode", "vibe")
|
||||
req.Header.Set("x-amzn-codewhisperer-optout", "true")
|
||||
req.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
|
||||
req.Header.Set("Amz-Sdk-Invocation-Id", uuid.NewString())
|
||||
|
||||
if account == nil {
|
||||
return
|
||||
}
|
||||
applyKiroConditionalHeaders(req, account)
|
||||
}
|
||||
|
||||
func accountProxyURL(account *Account) string {
|
||||
if account == nil || account.ProxyID == nil || account.Proxy == nil {
|
||||
return ""
|
||||
}
|
||||
return account.Proxy.URL()
|
||||
}
|
||||
|
||||
func kiroRuntimeEndpoint(region string) string {
|
||||
region = strings.TrimSpace(region)
|
||||
if region == "" {
|
||||
region = kiroDefaultRegion
|
||||
}
|
||||
switch region {
|
||||
case "us-east-1":
|
||||
return "https://q.us-east-1.amazonaws.com"
|
||||
case "eu-central-1":
|
||||
return "https://q.eu-central-1.amazonaws.com"
|
||||
case "us-gov-east-1":
|
||||
return "https://q-fips.us-gov-east-1.amazonaws.com"
|
||||
case "us-gov-west-1":
|
||||
return "https://q-fips.us-gov-west-1.amazonaws.com"
|
||||
case "us-iso-east-1":
|
||||
return "https://q.us-iso-east-1.c2s.ic.gov"
|
||||
case "us-isob-east-1":
|
||||
return "https://q.us-isob-east-1.sc2s.sgov.gov"
|
||||
case "us-isof-south-1":
|
||||
return "https://q.us-isof-south-1.csp.hci.ic.gov"
|
||||
case "us-isof-east-1":
|
||||
return "https://q.us-isof-east-1.csp.hci.ic.gov"
|
||||
default:
|
||||
if strings.HasPrefix(region, "us-gov-") {
|
||||
return "https://q-fips." + region + ".amazonaws.com"
|
||||
}
|
||||
return "https://q." + region + ".amazonaws.com"
|
||||
}
|
||||
}
|
||||
|
||||
func isLoopbackEndpoint(raw string) bool {
|
||||
parsed, err := url.Parse(strings.TrimSpace(raw))
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
host := strings.TrimSpace(parsed.Hostname())
|
||||
if host == "" {
|
||||
return false
|
||||
}
|
||||
if strings.EqualFold(host, "localhost") {
|
||||
return true
|
||||
}
|
||||
ip := net.ParseIP(host)
|
||||
return ip != nil && ip.IsLoopback()
|
||||
}
|
||||
|
||||
func mapKiroUsageToInfo(resp *kiroUsageLimitsResponse) *UsageInfo {
|
||||
now := time.Now()
|
||||
if resp == nil {
|
||||
return &UsageInfo{UpdatedAt: &now}
|
||||
}
|
||||
info := &UsageInfo{
|
||||
UpdatedAt: &now,
|
||||
KiroSubscriptionName: strings.TrimSpace(resp.SubscriptionInfo.SubscriptionTitle),
|
||||
KiroSubscriptionType: strings.TrimSpace(resp.SubscriptionInfo.Type),
|
||||
KiroOveragesEnabled: strings.EqualFold(strings.TrimSpace(resp.OverageConfiguration.OverageStatus), "ENABLED"),
|
||||
}
|
||||
|
||||
resetAt := parseKiroTimestamp(resp.NextDateReset)
|
||||
if credit := selectKiroCreditBreakdown(resp.UsageBreakdownList); credit != nil {
|
||||
if breakdownReset := parseKiroTimestamp(credit.NextDateReset); breakdownReset != nil {
|
||||
resetAt = breakdownReset
|
||||
}
|
||||
info.KiroCredit = &KiroCreditProgress{
|
||||
CurrentUsage: selectKiroFloat(credit.CurrentUsageWithPrecision, credit.CurrentUsage),
|
||||
UsageLimit: selectKiroFloat(credit.UsageLimitWithPrecision, credit.UsageLimit),
|
||||
PercentageUsed: percentageOrZero(selectKiroFloat(credit.CurrentUsageWithPrecision, credit.CurrentUsage), selectKiroFloat(credit.UsageLimitWithPrecision, credit.UsageLimit)),
|
||||
}
|
||||
info.KiroOverage = &KiroOverageInfo{
|
||||
CurrentOverages: selectKiroFloat(credit.CurrentOveragesWithPrecision, credit.CurrentOverages),
|
||||
OverageCharges: selectKiroFloat(credit.OverageCharges, nil),
|
||||
CurrencyCode: strings.TrimSpace(credit.Currency),
|
||||
CurrencySymbol: kiroCurrencySymbol(strings.TrimSpace(credit.Currency)),
|
||||
}
|
||||
if ft := credit.FreeTrialInfo; ft != nil && strings.EqualFold(strings.TrimSpace(ft.FreeTrialStatus), "ACTIVE") {
|
||||
expiry := parseKiroTimestamp(ft.FreeTrialExpiry)
|
||||
daysRemaining := 0
|
||||
if expiry != nil {
|
||||
daysRemaining = int(time.Until(*expiry).Hours() / 24)
|
||||
if time.Until(*expiry)%(24*time.Hour) != 0 {
|
||||
daysRemaining++
|
||||
}
|
||||
if daysRemaining < 0 {
|
||||
daysRemaining = 0
|
||||
}
|
||||
}
|
||||
current := selectKiroFloat(ft.CurrentUsageWithPrecision, ft.CurrentUsage)
|
||||
limit := selectKiroFloat(ft.UsageLimitWithPrecision, ft.UsageLimit)
|
||||
info.KiroBonus = &KiroCreditProgress{
|
||||
CurrentUsage: current,
|
||||
UsageLimit: limit,
|
||||
PercentageUsed: percentageOrZero(current, limit),
|
||||
DaysRemaining: daysRemaining,
|
||||
ExpiryDate: expiry,
|
||||
}
|
||||
}
|
||||
}
|
||||
info.KiroResetAt = resetAt
|
||||
setKiroQuotaStateFromUsage(info)
|
||||
return info
|
||||
}
|
||||
|
||||
func selectKiroCreditBreakdown(items []kiroUsageBreakdown) *kiroUsageBreakdown {
|
||||
for i := range items {
|
||||
if strings.EqualFold(strings.TrimSpace(items[i].ResourceType), "CREDIT") {
|
||||
return &items[i]
|
||||
}
|
||||
}
|
||||
if len(items) == 0 {
|
||||
return nil
|
||||
}
|
||||
return &items[0]
|
||||
}
|
||||
|
||||
func selectKiroFloat(preferred *float64, fallback *float64) float64 {
|
||||
switch {
|
||||
case preferred != nil:
|
||||
return *preferred
|
||||
case fallback != nil:
|
||||
return *fallback
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func percentageOrZero(current, limit float64) float64 {
|
||||
if limit <= 0 {
|
||||
return 0
|
||||
}
|
||||
return current / limit * 100
|
||||
}
|
||||
|
||||
func parseKiroTimestamp(raw any) *time.Time {
|
||||
if raw == nil {
|
||||
return nil
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case string:
|
||||
trimmed := strings.TrimSpace(v)
|
||||
if trimmed == "" {
|
||||
return nil
|
||||
}
|
||||
if parsed, err := time.Parse(time.RFC3339, trimmed); err == nil {
|
||||
return &parsed
|
||||
}
|
||||
if i, err := strconv.ParseInt(trimmed, 10, 64); err == nil {
|
||||
return unixishToTime(i)
|
||||
}
|
||||
if f, err := strconv.ParseFloat(trimmed, 64); err == nil {
|
||||
return unixishFloatToTime(f)
|
||||
}
|
||||
case float64:
|
||||
return unixishFloatToTime(v)
|
||||
case int64:
|
||||
return unixishToTime(v)
|
||||
case int:
|
||||
return unixishToTime(int64(v))
|
||||
case json.Number:
|
||||
if i, err := v.Int64(); err == nil {
|
||||
return unixishToTime(i)
|
||||
}
|
||||
if f, err := v.Float64(); err == nil {
|
||||
return unixishFloatToTime(f)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func unixishFloatToTime(v float64) *time.Time {
|
||||
if v <= 0 {
|
||||
return nil
|
||||
}
|
||||
if v >= 1e12 {
|
||||
t := time.UnixMilli(int64(v))
|
||||
return &t
|
||||
}
|
||||
t := time.Unix(int64(v), 0)
|
||||
return &t
|
||||
}
|
||||
|
||||
func unixishToTime(v int64) *time.Time {
|
||||
if v <= 0 {
|
||||
return nil
|
||||
}
|
||||
if v >= 1e12 {
|
||||
t := time.UnixMilli(v)
|
||||
return &t
|
||||
}
|
||||
t := time.Unix(v, 0)
|
||||
return &t
|
||||
}
|
||||
|
||||
func kiroCurrencySymbol(code string) string {
|
||||
switch strings.ToUpper(strings.TrimSpace(code)) {
|
||||
case "USD":
|
||||
return "$"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func buildKiroDegradedUsage(err error) *UsageInfo {
|
||||
now := time.Now()
|
||||
info := &UsageInfo{
|
||||
UpdatedAt: &now,
|
||||
Error: "usage API error",
|
||||
ErrorCode: errorCodeNetworkError,
|
||||
}
|
||||
if err == nil {
|
||||
return info
|
||||
}
|
||||
|
||||
info.Error = fmt.Sprintf("usage API error: %v", err)
|
||||
|
||||
classification := classifyKiroError(err)
|
||||
switch classification.Category {
|
||||
case kiroErrorAuthError:
|
||||
info.ErrorCode = errorCodeUnauthenticated
|
||||
info.NeedsReauth = true
|
||||
case kiroErrorRateLimited:
|
||||
info.ErrorCode = errorCodeRateLimited
|
||||
case kiroErrorQuotaExhausted:
|
||||
info.ErrorCode = errorCodeNetworkError
|
||||
info.KiroQuotaState = kiroQuotaStateCreditsExhausted
|
||||
info.KiroQuotaReason = classification.Message
|
||||
case kiroErrorOverageExhausted:
|
||||
info.ErrorCode = errorCodeNetworkError
|
||||
info.KiroQuotaState = kiroQuotaStateOverageExhausted
|
||||
info.KiroQuotaReason = classification.Message
|
||||
case kiroErrorSuspended, kiroErrorUsageForbidden, kiroErrorProfileError:
|
||||
info.ErrorCode = errorCodeForbidden
|
||||
default:
|
||||
info.ErrorCode = errorCodeNetworkError
|
||||
}
|
||||
return info
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) attachKiroRuntimeState(ctx context.Context, account *Account, usage *UsageInfo) {
|
||||
if s == nil || usage == nil || account == nil || account.Platform != PlatformKiro || s.kiroCooldownStore == nil {
|
||||
return
|
||||
}
|
||||
usage.KiroRuntimeState = ""
|
||||
usage.KiroRuntimeReason = ""
|
||||
usage.KiroRuntimeResetAt = nil
|
||||
state, err := s.kiroCooldownStore.GetState(ctx, buildKiroAccountKey(account))
|
||||
if err != nil || state == nil {
|
||||
return
|
||||
}
|
||||
usage.KiroRuntimeState, usage.KiroRuntimeReason, usage.KiroRuntimeResetAt = kiroRuntimeStateSnapshot(state)
|
||||
}
|
||||
|
||||
func (s *AccountUsageService) EnrichAccountWithKiroRuntimeState(ctx context.Context, account *Account) {
|
||||
if s == nil || account == nil || account.Platform != PlatformKiro || account.Type != AccountTypeOAuth {
|
||||
return
|
||||
}
|
||||
account.KiroQuotaState = ""
|
||||
account.KiroQuotaReason = ""
|
||||
account.KiroQuotaResetAt = nil
|
||||
account.KiroRuntimeState = ""
|
||||
account.KiroRuntimeReason = ""
|
||||
account.KiroRuntimeResetAt = nil
|
||||
if usage, ok := s.getCachedKiroUsage(account.ID); ok {
|
||||
account.KiroQuotaState = usage.KiroQuotaState
|
||||
account.KiroQuotaReason = usage.KiroQuotaReason
|
||||
account.KiroQuotaResetAt = usage.KiroQuotaResetAt
|
||||
}
|
||||
if s.kiroCooldownStore == nil {
|
||||
return
|
||||
}
|
||||
state, err := s.kiroCooldownStore.GetState(ctx, buildKiroAccountKey(account))
|
||||
if err != nil || state == nil {
|
||||
return
|
||||
}
|
||||
account.KiroRuntimeState, account.KiroRuntimeReason, account.KiroRuntimeResetAt = kiroRuntimeStateSnapshot(state)
|
||||
}
|
||||
|
||||
func setKiroQuotaStateFromUsage(info *UsageInfo) {
|
||||
if info == nil {
|
||||
return
|
||||
}
|
||||
info.KiroQuotaState = ""
|
||||
info.KiroQuotaReason = ""
|
||||
info.KiroQuotaResetAt = nil
|
||||
|
||||
creditExhausted := false
|
||||
if info.KiroCredit != nil && info.KiroCredit.UsageLimit > 0 {
|
||||
creditExhausted = info.KiroCredit.CurrentUsage >= info.KiroCredit.UsageLimit
|
||||
}
|
||||
overageActive := info.KiroOverage != nil &&
|
||||
(info.KiroOverage.CurrentOverages > 0 || info.KiroOverage.OverageCharges > 0)
|
||||
|
||||
switch {
|
||||
case info.KiroOveragesEnabled && (overageActive || creditExhausted):
|
||||
info.KiroQuotaState = kiroQuotaStateOverageActive
|
||||
info.KiroQuotaReason = "overages_enabled"
|
||||
info.KiroQuotaResetAt = info.KiroResetAt
|
||||
case creditExhausted:
|
||||
info.KiroQuotaState = kiroQuotaStateCreditsExhausted
|
||||
info.KiroQuotaReason = "credits_exhausted"
|
||||
info.KiroQuotaResetAt = info.KiroResetAt
|
||||
default:
|
||||
info.KiroQuotaState = kiroQuotaStateNormal
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user