feat(backend): add kiro account support

This commit is contained in:
nianzs
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)
}
+6
View File
@@ -221,6 +221,12 @@ func AccountFromServiceShallow(a *service.Account) *Account {
OverloadUntil: a.OverloadUntil,
TempUnschedulableUntil: a.TempUnschedulableUntil,
TempUnschedulableReason: a.TempUnschedulableReason,
KiroQuotaState: a.KiroQuotaState,
KiroQuotaReason: a.KiroQuotaReason,
KiroQuotaResetAt: a.KiroQuotaResetAt,
KiroRuntimeState: a.KiroRuntimeState,
KiroRuntimeReason: a.KiroRuntimeReason,
KiroRuntimeResetAt: a.KiroRuntimeResetAt,
SessionWindowStart: a.SessionWindowStart,
SessionWindowEnd: a.SessionWindowEnd,
SessionWindowStatus: a.SessionWindowStatus,
+6
View File
@@ -174,6 +174,12 @@ type Account struct {
TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until"`
TempUnschedulableReason string `json:"temp_unschedulable_reason"`
KiroQuotaState string `json:"kiro_quota_state,omitempty"`
KiroQuotaReason string `json:"kiro_quota_reason,omitempty"`
KiroQuotaResetAt *time.Time `json:"kiro_quota_reset_at,omitempty"`
KiroRuntimeState string `json:"kiro_runtime_state,omitempty"`
KiroRuntimeReason string `json:"kiro_runtime_reason,omitempty"`
KiroRuntimeResetAt *time.Time `json:"kiro_runtime_reset_at,omitempty"`
SessionWindowStart *time.Time `json:"session_window_start"`
SessionWindowEnd *time.Time `json:"session_window_end"`
+1
View File
@@ -17,6 +17,7 @@ type AdminHandlers struct {
OpenAIOAuth *admin.OpenAIOAuthHandler
GeminiOAuth *admin.GeminiOAuthHandler
AntigravityOAuth *admin.AntigravityOAuthHandler
KiroOAuth *admin.KiroOAuthHandler
Proxy *admin.ProxyHandler
Redeem *admin.RedeemHandler
Promo *admin.PromoHandler
+3
View File
@@ -20,6 +20,7 @@ func ProvideAdminHandlers(
openaiOAuthHandler *admin.OpenAIOAuthHandler,
geminiOAuthHandler *admin.GeminiOAuthHandler,
antigravityOAuthHandler *admin.AntigravityOAuthHandler,
kiroOAuthHandler *admin.KiroOAuthHandler,
proxyHandler *admin.ProxyHandler,
redeemHandler *admin.RedeemHandler,
promoHandler *admin.PromoHandler,
@@ -51,6 +52,7 @@ func ProvideAdminHandlers(
OpenAIOAuth: openaiOAuthHandler,
GeminiOAuth: geminiOAuthHandler,
AntigravityOAuth: antigravityOAuthHandler,
KiroOAuth: kiroOAuthHandler,
Proxy: proxyHandler,
Redeem: redeemHandler,
Promo: promoHandler,
@@ -154,6 +156,7 @@ var ProviderSet = wire.NewSet(
admin.NewOpenAIOAuthHandler,
admin.NewGeminiOAuthHandler,
admin.NewAntigravityOAuthHandler,
admin.NewKiroOAuthHandler,
admin.NewProxyHandler,
admin.NewRedeemHandler,
admin.NewPromoHandler,