diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index 28902fe55e..f92f452df2 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -400,6 +400,9 @@ func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth) gin.H { "source": "memory", "size": int64(0), } + entry["success"] = auth.Success + entry["failed"] = auth.Failed + entry["recent_requests"] = auth.RecentRequestsSnapshot(time.Now()) if email := authEmail(auth); email != "" { entry["email"] = email } @@ -443,6 +446,9 @@ func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth) gin.H { if claims := extractCodexIDTokenClaims(auth); claims != nil { entry["id_token"] = claims } + if kiroQuota := h.getKiroQuotaCached(auth); kiroQuota != nil { + entry["kiro_quota"] = kiroQuota + } // Expose priority from Attributes (set by synthesizer from JSON "priority" field). // Fall back to Metadata for auths registered via UploadAuthFile (no synthesizer). if p := strings.TrimSpace(authAttribute(auth, "priority")); p != "" { diff --git a/internal/api/handlers/management/kiro_quota.go b/internal/api/handlers/management/kiro_quota.go new file mode 100644 index 0000000000..421dfa77ec --- /dev/null +++ b/internal/api/handlers/management/kiro_quota.go @@ -0,0 +1,174 @@ +package management + +import ( + "context" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + log "github.com/sirupsen/logrus" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro" + "github.com/router-for-me/CLIProxyAPI/v6/internal/util" + coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" +) + +// GetKiroQuota fetches Kiro (AWS CodeWhisperer) usage quota information. +// +// Endpoint: +// +// GET /v0/management/kiro-quota +// +// Query Parameters (optional): +// - auth_index: The credential "auth_index" from GET /v0/management/auth-files. +// If omitted, uses the first available Kiro credential. +// +// Response: +// +// Returns the UsageQuotaResponse with usage breakdown and subscription info. +// +// Example: +// +// curl -sS -X GET "http://127.0.0.1:8317/v0/management/kiro-quota?auth_index=" \ +// -H "Authorization: Bearer " +func (h *Handler) GetKiroQuota(c *gin.Context) { + authIndex := strings.TrimSpace(c.Query("auth_index")) + if authIndex == "" { + authIndex = strings.TrimSpace(c.Query("authIndex")) + } + if authIndex == "" { + authIndex = strings.TrimSpace(c.Query("AuthIndex")) + } + + auth := h.findKiroAuth(authIndex) + if auth == nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "no kiro credential found"}) + return + } + + // Extract token data from auth metadata + tokenData := extractKiroTokenData(auth) + if tokenData == nil || tokenData.AccessToken == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "kiro access token not available (token may need refresh)"}) + return + } + + // Create usage checker with proxy-aware HTTP client + checker := kiro.NewUsageCheckerWithClient( + util.SetProxy(&h.cfg.SDKConfig, &http.Client{Timeout: 30 * time.Second}), + ) + + ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second) + defer cancel() + + usage, err := checker.CheckUsage(ctx, tokenData) + if err != nil { + log.WithError(err).Debug("kiro quota request failed") + c.JSON(http.StatusBadGateway, gin.H{"error": "kiro quota request failed: " + err.Error()}) + return + } + + // Build enriched response + response := gin.H{ + "usage": usage, + "quota_status": buildKiroQuotaStatus(usage), + "auth_index": auth.Index, + "auth_name": auth.FileName, + } + + c.JSON(http.StatusOK, response) +} + +// findKiroAuth locates a Kiro credential by auth_index or returns the first available one. +func (h *Handler) findKiroAuth(authIndex string) *coreauth.Auth { + if h == nil || h.authManager == nil { + return nil + } + + auths := h.authManager.List() + var firstKiro *coreauth.Auth + + for _, auth := range auths { + if auth == nil { + continue + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if provider != "kiro" { + continue + } + if auth.Disabled { + continue + } + if firstKiro == nil { + firstKiro = auth + } + if authIndex != "" { + auth.EnsureIndex() + if auth.Index == authIndex { + return auth + } + } + } + + if authIndex == "" { + return firstKiro + } + return nil +} + +// extractKiroTokenData extracts KiroTokenData from a coreauth.Auth's Metadata. +func extractKiroTokenData(auth *coreauth.Auth) *kiro.KiroTokenData { + if auth == nil || auth.Metadata == nil { + return nil + } + + accessToken, _ := auth.Metadata["access_token"].(string) + refreshToken, _ := auth.Metadata["refresh_token"].(string) + profileArn, _ := auth.Metadata["profile_arn"].(string) + clientID, _ := auth.Metadata["client_id"].(string) + clientSecret, _ := auth.Metadata["client_secret"].(string) + region, _ := auth.Metadata["region"].(string) + startURL, _ := auth.Metadata["start_url"].(string) + + if accessToken == "" { + return nil + } + + return &kiro.KiroTokenData{ + AccessToken: accessToken, + RefreshToken: refreshToken, + ProfileArn: profileArn, + ClientID: clientID, + ClientSecret: clientSecret, + Region: region, + StartURL: startURL, + } +} + +// buildKiroQuotaStatus builds a summary status from the usage response. +func buildKiroQuotaStatus(usage *kiro.UsageQuotaResponse) gin.H { + if usage == nil { + return gin.H{"exhausted": true, "remaining": 0} + } + + remaining := kiro.GetRemainingQuota(usage) + exhausted := kiro.IsQuotaExhausted(usage) + percentage := kiro.GetUsagePercentage(usage) + + status := gin.H{ + "exhausted": exhausted, + "remaining": remaining, + "usage_percentage": percentage, + } + + if usage.NextDateReset > 0 { + status["next_reset"] = time.Unix(int64(usage.NextDateReset/1000), 0) + } + + if usage.SubscriptionInfo != nil { + status["subscription"] = usage.SubscriptionInfo + } + + return status +} diff --git a/internal/api/handlers/management/kiro_quota_cache.go b/internal/api/handlers/management/kiro_quota_cache.go new file mode 100644 index 0000000000..c430ed3165 --- /dev/null +++ b/internal/api/handlers/management/kiro_quota_cache.go @@ -0,0 +1,141 @@ +package management + +import ( + "context" + "net/http" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + log "github.com/sirupsen/logrus" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro" + "github.com/router-for-me/CLIProxyAPI/v6/internal/util" + coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" +) + +// kiroQuotaCache stores cached quota info for Kiro auth entries. +var ( + kiroQuotaMu sync.RWMutex + kiroQuotaStore = make(map[string]gin.H) // keyed by auth ID +) + +// getKiroQuotaCached returns cached quota info for a Kiro auth entry. +func (h *Handler) getKiroQuotaCached(auth *coreauth.Auth) gin.H { + if auth == nil { + return nil + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if provider != "kiro" { + return nil + } + + kiroQuotaMu.RLock() + result, ok := kiroQuotaStore[auth.ID] + kiroQuotaMu.RUnlock() + + if ok { + return result + } + + // If not cached, try to fetch synchronously (first time only) + return h.fetchAndCacheKiroQuota(auth) +} + +// fetchAndCacheKiroQuota fetches quota for a single Kiro auth and caches it. +func (h *Handler) fetchAndCacheKiroQuota(auth *coreauth.Auth) gin.H { + if auth == nil || auth.Metadata == nil { + return nil + } + + tokenData := extractKiroTokenData(auth) + if tokenData == nil || tokenData.AccessToken == "" { + return nil + } + + checker := kiro.NewUsageCheckerWithClient( + util.SetProxy(&h.cfg.SDKConfig, &http.Client{Timeout: 15 * time.Second}), + ) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + usage, err := checker.CheckUsage(ctx, tokenData) + if err != nil { + log.WithError(err).Debugf("kiro quota fetch failed for %s", auth.ID) + return nil + } + + result := buildKiroQuotaEntry(usage) + + kiroQuotaMu.Lock() + kiroQuotaStore[auth.ID] = result + kiroQuotaMu.Unlock() + + return result +} + +// buildKiroQuotaEntry builds the quota info map from a usage response. +func buildKiroQuotaEntry(usage *kiro.UsageQuotaResponse) gin.H { + if usage == nil || len(usage.UsageBreakdownList) == 0 { + return nil + } + + bd := usage.UsageBreakdownList[0] + result := gin.H{ + "resource_type": bd.ResourceType, + "used": bd.CurrentUsageWithPrecision, + "limit": bd.UsageLimitWithPrecision, + "remaining": bd.UsageLimitWithPrecision - bd.CurrentUsageWithPrecision, + "usage_percentage": kiro.GetUsagePercentage(usage), + "exhausted": kiro.IsQuotaExhausted(usage), + } + + if usage.SubscriptionInfo != nil { + result["plan"] = usage.SubscriptionInfo.SubscriptionTitle + } + + if usage.NextDateReset > 0 { + result["next_reset"] = time.Unix(int64(usage.NextDateReset/1000), 0) + } + + return result +} + +// StartKiroQuotaRefresher starts a background goroutine that periodically +// refreshes Kiro quota info for all active Kiro auth entries. +func (h *Handler) StartKiroQuotaRefresher() { + go func() { + // Initial delay to let the service start up + time.Sleep(10 * time.Second) + h.refreshAllKiroQuotas() + + ticker := time.NewTicker(5 * time.Minute) + defer ticker.Stop() + for range ticker.C { + h.refreshAllKiroQuotas() + } + }() +} + +// refreshAllKiroQuotas fetches quota for all active Kiro auth entries. +func (h *Handler) refreshAllKiroQuotas() { + if h == nil || h.authManager == nil { + return + } + + auths := h.authManager.List() + for _, auth := range auths { + if auth == nil || auth.Disabled { + continue + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if provider != "kiro" { + continue + } + h.fetchAndCacheKiroQuota(auth) + // Small delay between requests to avoid rate limiting + time.Sleep(2 * time.Second) + } +} diff --git a/internal/api/server.go b/internal/api/server.go index 8fd7f4680f..07832d9619 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -665,6 +665,8 @@ func (s *Server) registerManagementRoutes() { mgmt.PATCH("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel) mgmt.GET("/copilot-quota", s.mgmt.GetCopilotQuota) + mgmt.GET("/kiro-quota", s.mgmt.GetKiroQuota) + s.mgmt.StartKiroQuotaRefresher() mgmt.GET("/api-keys", s.mgmt.GetAPIKeys) mgmt.PUT("/api-keys", s.mgmt.PutAPIKeys) diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 6628c20b8a..0e3bb71581 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1121,6 +1121,9 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { auth.Index = existing.Index auth.indexAssigned = existing.indexAssigned } + auth.Success = existing.Success + auth.Failed = existing.Failed + auth.recentRequests = existing.recentRequests if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { auth.ModelStates = existing.ModelStates @@ -1977,6 +1980,13 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { if auth, ok := m.auths[result.AuthID]; ok && auth != nil { now := time.Now() + auth.recordRecentRequest(now, result.Success) + if result.Success { + auth.Success++ + } else { + auth.Failed++ + } + if result.Success { if result.Model != "" { state := ensureModelState(auth, result.Model) diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index 16814eb899..f3ffcbf501 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -90,8 +90,15 @@ type Auth struct { ModelStates map[string]*ModelState `json:"model_states,omitempty"` // Runtime carries non-serialisable data used during execution (in-memory only). + // Success counts the total number of successful requests routed through this auth. + Success int64 `json:"-"` + // Failed counts the total number of failed requests routed through this auth. + Failed int64 `json:"-"` + Runtime any `json:"-"` + recentRequests recentRequestRing `json:"-"` + indexAssigned bool `json:"-"` } @@ -583,3 +590,92 @@ func normaliseUnix(raw int64) time.Time { } return time.Unix(raw, 0) } + +// --- Recent Requests Health Tracking --- + +const ( + recentRequestBucketSeconds int64 = 10 * 60 + recentRequestBucketCount = 20 +) + +type recentRequestBucket struct { + bucketID int64 + success int64 + failed int64 +} + +type recentRequestRing struct { + buckets [recentRequestBucketCount]recentRequestBucket +} + +// RecentRequestBucket is the exported per-bucket snapshot returned by the management API. +type RecentRequestBucket struct { + Time string `json:"time"` + Success int64 `json:"success"` + Failed int64 `json:"failed"` +} + +func recentRequestBucketID(now time.Time) int64 { + if now.IsZero() { + return 0 + } + return now.Unix() / recentRequestBucketSeconds +} + +func recentRequestBucketIndex(bucketID int64) int { + mod := bucketID % int64(recentRequestBucketCount) + if mod < 0 { + mod += int64(recentRequestBucketCount) + } + return int(mod) +} + +func formatRecentRequestBucketLabel(bucketID int64) string { + start := time.Unix(bucketID*recentRequestBucketSeconds, 0).In(time.Local) + end := start.Add(time.Duration(recentRequestBucketSeconds) * time.Second) + return start.Format("15:04") + "-" + end.Format("15:04") +} + +func (a *Auth) recordRecentRequest(now time.Time, success bool) { + if a == nil { + return + } + bucketID := recentRequestBucketID(now) + idx := recentRequestBucketIndex(bucketID) + bucket := &a.recentRequests.buckets[idx] + if bucket.bucketID != bucketID { + bucket.bucketID = bucketID + bucket.success = 0 + bucket.failed = 0 + } + if success { + bucket.success++ + } else { + bucket.failed++ + } +} + +// RecentRequestsSnapshot returns the last 20 ten-minute buckets of request counts. +func (a *Auth) RecentRequestsSnapshot(now time.Time) []RecentRequestBucket { + out := make([]RecentRequestBucket, 0, recentRequestBucketCount) + if a == nil { + return out + } + + currentBucketID := recentRequestBucketID(now) + for i := recentRequestBucketCount - 1; i >= 0; i-- { + bucketID := currentBucketID - int64(i) + idx := recentRequestBucketIndex(bucketID) + bucket := a.recentRequests.buckets[idx] + entry := RecentRequestBucket{ + Time: formatRecentRequestBucketLabel(bucketID), + } + if bucket.bucketID == bucketID { + entry.Success = bucket.success + entry.Failed = bucket.failed + } + out = append(out, entry) + } + + return out +}