feat: Add validation and account management functionality (#21)

* feat: Add validation and account management functionality - Add validation for clientID and clientSecret in refreshOIDCToken function - Add weight field for load balancing priority in Account struct - Implement weighted轮询策略以根据账号权重分配选择概率。 - Add batch account management functionality including enabling, disabling, refreshing, and retrieving account details. - Update Kiro API version and adjust user agent strings to reflect new version numbers. - Update Kiro version and modify user agent strings and header settings. - Refactor model mapping to an ordered list for precise key matching. - Add account bulk actions and filtering toolbar to index.html * feat: Add logic to skip accounts with exhausted usage limits - Add logic to skip accounts with exhausted usage limits when selecting the next account.
2026-02-23 21:47:17 +08:00
parent d71bf09dde
commit ad7aabd554
7 changed files with 323 additions and 23 deletions
--- a/proxy/handler.go
+++ b/proxy/handler.go
@@ -1540,6 +1540,8 @@ func (h *Handler) handleAdminAPI(w http.ResponseWriter, r *http.Request) {
 		h.apiGetAccounts(w, r)
 	case path == "/accounts" && r.Method == "POST":
 		h.apiAddAccount(w, r)
+	case path == "/accounts/batch" && r.Method == "POST":
+		h.apiBatchAccounts(w, r)
 	case strings.HasPrefix(path, "/accounts/") && strings.HasSuffix(path, "/refresh") && r.Method == "POST":
 		id := strings.TrimSuffix(strings.TrimPrefix(path, "/accounts/"), "/refresh")
 		h.apiRefreshAccount(w, r, id)
@@ -1626,6 +1628,7 @@ func (h *Handler) apiGetAccounts(w http.ResponseWriter, r *http.Request) {
 			"expiresAt":         a.ExpiresAt,
 			"hasToken":          a.AccessToken != "",
 			"machineId":         a.MachineId,
+			"weight":            a.Weight,
 			"subscriptionType":  a.SubscriptionType,
 			"subscriptionTitle": a.SubscriptionTitle,
 			"daysRemaining":     a.DaysRemaining,
@@ -1717,6 +1720,9 @@ func (h *Handler) apiUpdateAccount(w http.ResponseWriter, r *http.Request, id st
 	if v, ok := updates["machineId"].(string); ok {
 		existing.MachineId = v
 	}
+	if v, ok := updates["weight"].(float64); ok {
+		existing.Weight = int(v)
+	}

 	if err := config.UpdateAccount(id, *existing); err != nil {
 		w.WriteHeader(500)
@@ -1728,6 +1734,95 @@ func (h *Handler) apiUpdateAccount(w http.ResponseWriter, r *http.Request, id st
 	json.NewEncoder(w).Encode(map[string]bool{"success": true})
 }

+// apiBatchAccounts 批量操作账号（启用/禁用/刷新）
+func (h *Handler) apiBatchAccounts(w http.ResponseWriter, r *http.Request) {
+	var req struct {
+		IDs    []string `json:"ids"`
+		Action string   `json:"action"` // "enable", "disable", "refresh"
+	}
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		w.WriteHeader(400)
+		json.NewEncoder(w).Encode(map[string]string{"error": "Invalid JSON"})
+		return
+	}
+	if len(req.IDs) == 0 {
+		w.WriteHeader(400)
+		json.NewEncoder(w).Encode(map[string]string{"error": "No account IDs provided"})
+		return
+	}
+
+	switch req.Action {
+	case "enable", "disable":
+		enabled := req.Action == "enable"
+		accounts := config.GetAccounts()
+		idSet := make(map[string]bool)
+		for _, id := range req.IDs {
+			idSet[id] = true
+		}
+		for _, a := range accounts {
+			if idSet[a.ID] {
+				a.Enabled = enabled
+				if enabled && a.BanStatus != "" && a.BanStatus != "ACTIVE" {
+					a.BanStatus = "ACTIVE"
+					a.BanReason = ""
+					a.BanTime = 0
+				}
+				config.UpdateAccount(a.ID, a)
+			}
+		}
+		h.pool.Reload()
+		json.NewEncoder(w).Encode(map[string]interface{}{"success": true, "count": len(req.IDs)})
+
+	case "refresh":
+		successCount := 0
+		failCount := 0
+		for _, id := range req.IDs {
+			accounts := config.GetAccounts()
+			var account *config.Account
+			for i := range accounts {
+				if accounts[i].ID == id {
+					account = &accounts[i]
+					break
+				}
+			}
+			if account == nil {
+				failCount++
+				continue
+			}
+			// 刷新 token
+			if account.RefreshToken != "" {
+				if newAccess, newRefresh, newExpires, err := auth.RefreshToken(account); err == nil {
+					account.AccessToken = newAccess
+					if newRefresh != "" {
+						account.RefreshToken = newRefresh
+					}
+					account.ExpiresAt = newExpires
+					config.UpdateAccountToken(id, newAccess, newRefresh, newExpires)
+					h.pool.UpdateToken(id, newAccess, newRefresh, newExpires)
+				}
+			}
+			// 刷新账户信息
+			info, err := RefreshAccountInfo(account)
+			if err != nil {
+				failCount++
+				continue
+			}
+			config.UpdateAccountInfo(id, *info)
+			successCount++
+		}
+		h.pool.Reload()
+		json.NewEncoder(w).Encode(map[string]interface{}{
+			"success": true,
+			"refreshed": successCount,
+			"failed":    failCount,
+		})
+
+	default:
+		w.WriteHeader(400)
+		json.NewEncoder(w).Encode(map[string]string{"error": "Invalid action: " + req.Action})
+	}
+}
+
 func (h *Handler) apiStartIamSso(w http.ResponseWriter, r *http.Request) {
 	var req struct {
 		StartUrl string `json:"startUrl"`
--- a/proxy/kiro.go
+++ b/proxy/kiro.go
@@ -16,7 +16,7 @@ import (
 	"github.com/google/uuid"
 )

-const KiroVersion = "0.6.18"
+const KiroVersion = "0.7.45"

 // 双端点配置（429 时自动 fallback）
 type kiroEndpoint struct {
@@ -168,11 +168,11 @@ func CallKiroAPI(account *config.Account, payload *KiroPayload, callback *KiroSt
 	machineId := account.MachineId
 	var userAgent, amzUserAgent string
 	if machineId != "" {
-		userAgent = fmt.Sprintf("aws-sdk-js/1.0.18 ua/2.1 os/linux lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.18 m/E KiroIDE-%s-%s", KiroVersion, machineId)
-		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.18 KiroIDE %s %s", KiroVersion, machineId)
+		userAgent = fmt.Sprintf("aws-sdk-js/1.0.27 ua/2.1 os/linux lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E KiroIDE-%s-%s", KiroVersion, machineId)
+		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.27 KiroIDE %s %s", KiroVersion, machineId)
 	} else {
-		userAgent = fmt.Sprintf("aws-sdk-js/1.0.18 ua/2.1 os/linux lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.18 m/E KiroIDE-%s", KiroVersion)
-		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.18 KiroIDE %s", KiroVersion)
+		userAgent = fmt.Sprintf("aws-sdk-js/1.0.27 ua/2.1 os/linux lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E KiroIDE-%s", KiroVersion)
+		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.27 KiroIDE %s", KiroVersion)
 	}

 	// 根据配置排序端点
@@ -195,7 +195,7 @@ func CallKiroAPI(account *config.Account, payload *KiroPayload, callback *KiroSt
 		req.Header.Set("X-Amz-Target", ep.AmzTarget)
 		req.Header.Set("User-Agent", userAgent)
 		req.Header.Set("X-Amz-User-Agent", amzUserAgent)
-		req.Header.Set("x-amzn-kiro-agent-mode", "spec")
+		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.New().String())
--- a/proxy/kiro_api.go
+++ b/proxy/kiro_api.go
@@ -12,7 +12,7 @@ import (

 const (
 	kiroRestAPIBase = "https://codewhisperer.us-east-1.amazonaws.com"
-	kiroVersion     = "0.6.18"
+	kiroVersion     = "0.7.45"
 )

 // GetUsageLimits 获取账户使用量和订阅信息
@@ -113,11 +113,11 @@ func setKiroHeaders(req *http.Request, account *config.Account) {
 	machineId := account.MachineId
 	var userAgent, amzUserAgent string
 	if machineId != "" {
-		userAgent = fmt.Sprintf("aws-sdk-js/1.0.18 ua/2.1 os/windows lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.18 m/E KiroIDE-%s-%s", kiroVersion, machineId)
-		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.18 KiroIDE %s %s", kiroVersion, machineId)
+		userAgent = fmt.Sprintf("aws-sdk-js/1.0.27 ua/2.1 os/linux lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E KiroIDE-%s-%s", kiroVersion, machineId)
+		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.27 KiroIDE %s %s", kiroVersion, machineId)
 	} else {
-		userAgent = fmt.Sprintf("aws-sdk-js/1.0.18 ua/2.1 os/windows lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.18 m/E KiroIDE-%s", kiroVersion)
-		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.18 KiroIDE-%s", kiroVersion)
+		userAgent = fmt.Sprintf("aws-sdk-js/1.0.27 ua/2.1 os/linux lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E KiroIDE-%s", kiroVersion)
+		amzUserAgent = fmt.Sprintf("aws-sdk-js/1.0.27 KiroIDE %s", kiroVersion)
 	}

 	req.Header.Set("Authorization", "Bearer "+account.AccessToken)