feat(backend): add kiro account support

2026-04-29 16:29:21 +08:00
parent 9d801595c9
commit 05bc424c9a
60 changed files with 11916 additions and 38 deletions
@@ -22,6 +22,7 @@ import (
 	"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
 	infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
 	"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
+	kiropkg "github.com/Wei-Shaw/sub2api/internal/pkg/kiro"
 	"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
 	"github.com/Wei-Shaw/sub2api/internal/pkg/response"
 	"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
@@ -179,6 +180,9 @@ type AccountWithConcurrency struct {
 const accountListGroupUngroupedQueryValue = "ungrouped"

 func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, account *service.Account) AccountWithConcurrency {
+	if h.accountUsageService != nil {
+		h.accountUsageService.EnrichAccountWithKiroRuntimeState(ctx, account)
+	}
 	item := AccountWithConcurrency{
 		Account:            dto.AccountFromService(account),
 		CurrentConcurrency: 0,
@@ -351,6 +355,9 @@ func (h *AccountHandler) List(c *gin.Context) {
 	result := make([]AccountWithConcurrency, len(accounts))
 	for i := range accounts {
 		acc := &accounts[i]
+		if h.accountUsageService != nil {
+			h.accountUsageService.EnrichAccountWithKiroRuntimeState(c.Request.Context(), acc)
+		}
 		item := AccountWithConcurrency{
 			Account:            dto.AccountFromService(acc),
 			CurrentConcurrency: concurrencyCounts[acc.ID],
@@ -1913,6 +1920,18 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
 		return
 	}

+	// Handle Kiro accounts
+	if account.Platform == service.PlatformKiro {
+		mapping := account.GetModelMapping()
+		if len(mapping) == 0 {
+			response.Success(c, kiropkg.DefaultModels)
+			return
+		}
+
+		response.Success(c, buildMappedKiroModels(mapping))
+		return
+	}
+
 	// Handle Claude/Anthropic accounts
 	// For OAuth and Setup-Token accounts: return default models
 	if account.IsOAuth() {
@@ -1954,6 +1973,28 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
 	response.Success(c, models)
 }

+func buildMappedKiroModels(mapping map[string]string) []kiropkg.Model {
+	models := make([]kiropkg.Model, 0, len(mapping))
+	for requestedModel := range mapping {
+		var found bool
+		for _, dm := range kiropkg.DefaultModels {
+			if dm.ID == requestedModel {
+				models = append(models, dm)
+				found = true
+				break
+			}
+		}
+		if !found {
+			models = append(models, kiropkg.Model{
+				ID:          requestedModel,
+				Type:        "model",
+				DisplayName: requestedModel,
+			})
+		}
+	}
+	return models
+}
+
 // SetPrivacy handles setting privacy for a single OpenAI/Antigravity OAuth account
 // POST /api/v1/admin/accounts/:id/set-privacy
 func (h *AccountHandler) SetPrivacy(c *gin.Context) {
@@ -2166,6 +2207,12 @@ func (h *AccountHandler) GetAntigravityDefaultModelMapping(c *gin.Context) {
 	response.Success(c, domain.DefaultAntigravityModelMapping)
 }

+// GetKiroDefaultModelMapping 获取 Kiro 平台的默认模型映射
+// GET /api/v1/admin/accounts/kiro/default-model-mapping
+func (h *AccountHandler) GetKiroDefaultModelMapping(c *gin.Context) {
+	response.Success(c, domain.DefaultKiroModelMapping)
+}
+
 // sanitizeExtraBaseRPM 对 extra map 中的 base_rpm 值进行范围校验和归一化。
 // 负值归零，超过 10000 截断为 10000。extra 为 nil 或不含 base_rpm 时无操作。
 func sanitizeExtraBaseRPM(extra map[string]any) {
@@ -0,0 +1,149 @@
+package admin
+
+import (
+	"github.com/Wei-Shaw/sub2api/internal/pkg/response"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/gin-gonic/gin"
+)
+
+type KiroOAuthHandler struct {
+	kiroOAuthService *service.KiroOAuthService
+}
+
+func NewKiroOAuthHandler(kiroOAuthService *service.KiroOAuthService) *KiroOAuthHandler {
+	return &KiroOAuthHandler{kiroOAuthService: kiroOAuthService}
+}
+
+type KiroGenerateAuthURLRequest struct {
+	ProxyID  *int64 `json:"proxy_id"`
+	Provider string `json:"provider"`
+}
+
+func (h *KiroOAuthHandler) GenerateAuthURL(c *gin.Context) {
+	var req KiroGenerateAuthURLRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "请求无效: "+err.Error())
+		return
+	}
+	result, err := h.kiroOAuthService.GenerateAuthURL(c.Request.Context(), &service.KiroGenerateAuthURLInput{
+		ProxyID:  req.ProxyID,
+		Provider: req.Provider,
+	})
+	if err != nil {
+		response.BadRequest(c, "生成授权链接失败: "+err.Error())
+		return
+	}
+	response.Success(c, result)
+}
+
+type KiroGenerateIDCAuthURLRequest struct {
+	ProxyID  *int64 `json:"proxy_id"`
+	StartURL string `json:"start_url"`
+	Region   string `json:"region"`
+}
+
+func (h *KiroOAuthHandler) GenerateIDCAuthURL(c *gin.Context) {
+	var req KiroGenerateIDCAuthURLRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "请求无效: "+err.Error())
+		return
+	}
+	result, err := h.kiroOAuthService.GenerateIDCAuthURL(c.Request.Context(), &service.KiroGenerateIDCAuthURLInput{
+		ProxyID:  req.ProxyID,
+		StartURL: req.StartURL,
+		Region:   req.Region,
+	})
+	if err != nil {
+		response.BadRequest(c, "生成 IDC 授权链接失败: "+err.Error())
+		return
+	}
+	response.Success(c, result)
+}
+
+type KiroExchangeCodeRequest struct {
+	SessionID    string `json:"session_id" binding:"required"`
+	State        string `json:"state" binding:"required"`
+	Code         string `json:"code" binding:"required"`
+	CallbackPath string `json:"callback_path"`
+	LoginOption  string `json:"login_option"`
+	ProxyID      *int64 `json:"proxy_id"`
+}
+
+func (h *KiroOAuthHandler) ExchangeCode(c *gin.Context) {
+	var req KiroExchangeCodeRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "请求无效: "+err.Error())
+		return
+	}
+	tokenInfo, err := h.kiroOAuthService.ExchangeCode(c.Request.Context(), &service.KiroExchangeCodeInput{
+		SessionID:    req.SessionID,
+		State:        req.State,
+		Code:         req.Code,
+		CallbackPath: req.CallbackPath,
+		LoginOption:  req.LoginOption,
+		ProxyID:      req.ProxyID,
+	})
+	if err != nil {
+		response.BadRequest(c, "Token 交换失败: "+err.Error())
+		return
+	}
+	response.Success(c, tokenInfo)
+}
+
+type KiroRefreshTokenRequest struct {
+	RefreshToken string `json:"refresh_token" binding:"required"`
+	AuthMethod   string `json:"auth_method"`
+	Provider     string `json:"provider"`
+	ClientID     string `json:"client_id"`
+	ClientSecret string `json:"client_secret"`
+	StartURL     string `json:"start_url"`
+	Region       string `json:"region"`
+	ProfileArn   string `json:"profile_arn"`
+	ProxyID      *int64 `json:"proxy_id"`
+}
+
+func (h *KiroOAuthHandler) RefreshToken(c *gin.Context) {
+	var req KiroRefreshTokenRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "请求无效: "+err.Error())
+		return
+	}
+	tokenInfo, err := h.kiroOAuthService.RefreshToken(c.Request.Context(), &service.KiroRefreshTokenInput{
+		RefreshToken: req.RefreshToken,
+		AuthMethod:   req.AuthMethod,
+		Provider:     req.Provider,
+		ClientID:     req.ClientID,
+		ClientSecret: req.ClientSecret,
+		StartURL:     req.StartURL,
+		Region:       req.Region,
+		ProfileArn:   req.ProfileArn,
+		ProxyID:      req.ProxyID,
+	})
+	if err != nil {
+		response.BadRequest(c, "刷新 Kiro Token 失败: "+err.Error())
+		return
+	}
+	response.Success(c, tokenInfo)
+}
+
+type KiroImportTokenRequest struct {
+	TokenJSON              string `json:"token_json" binding:"required"`
+	DeviceRegistrationJSON string `json:"device_registration_json"`
+}
+
+func (h *KiroOAuthHandler) ImportToken(c *gin.Context) {
+	var req KiroImportTokenRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "请求无效: "+err.Error())
+		return
+	}
+	tokenInfo, err := h.kiroOAuthService.ImportToken(&service.KiroImportTokenInput{
+		TokenJSON:              req.TokenJSON,
+		DeviceRegistrationJSON: req.DeviceRegistrationJSON,
+	})
+	if err != nil {
+		response.BadRequest(c, "导入 Kiro Token 失败: "+err.Error())
+		return
+	}
+	response.Success(c, tokenInfo)
+}