feat(backend): add kiro account support

2026-04-30 14:02:05 +08:00
parent 9d801595c9
commit 05bc424c9a
60 changed files with 11916 additions and 38 deletions
@@ -20,6 +20,7 @@ import (
 	"github.com/Wei-Shaw/sub2api/internal/config"
 	"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
 	"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/util/urlvalidator"
 	"github.com/gin-gonic/gin"
@@ -65,6 +66,7 @@ type AccountTestService struct {
 	accountRepo               AccountRepository
 	geminiTokenProvider       *GeminiTokenProvider
 	claudeTokenProvider       *ClaudeTokenProvider
+	kiroTokenProvider         *KiroTokenProvider
 	antigravityGatewayService *AntigravityGatewayService
 	httpUpstream              HTTPUpstream
 	cfg                       *config.Config
@@ -76,6 +78,7 @@ func NewAccountTestService(
 	accountRepo AccountRepository,
 	geminiTokenProvider *GeminiTokenProvider,
 	claudeTokenProvider *ClaudeTokenProvider,
+	kiroTokenProvider *KiroTokenProvider,
 	antigravityGatewayService *AntigravityGatewayService,
 	httpUpstream HTTPUpstream,
 	cfg *config.Config,
@@ -85,6 +88,7 @@ func NewAccountTestService(
 		accountRepo:               accountRepo,
 		geminiTokenProvider:       geminiTokenProvider,
 		claudeTokenProvider:       claudeTokenProvider,
+		kiroTokenProvider:         kiroTokenProvider,
 		antigravityGatewayService: antigravityGatewayService,
 		httpUpstream:              httpUpstream,
 		cfg:                       cfg,
@@ -191,6 +195,10 @@ func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int
 		return s.routeAntigravityTest(c, account, modelID, prompt)
 	}

+	if account.IsKiro() && account.Type == AccountTypeOAuth {
+		return s.testKiroAccountConnection(c, account, modelID)
+	}
+
 	return s.testClaudeAccountConnection(c, account, modelID)
 }

@@ -239,6 +247,9 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
 		}

 		baseURL := account.GetBaseURL()
+		if baseURL == "" && account.Platform == PlatformKiro {
+			return s.sendErrorAndEnd(c, "Kiro API Key accounts require a Base URL")
+		}
 		if baseURL == "" {
 			baseURL = "https://api.anthropic.com"
 		}
@@ -387,6 +398,149 @@ func (s *AccountTestService) testClaudeVertexServiceAccountConnection(c *gin.Con
 	return s.processClaudeStream(c, resp.Body)
 }

+func (s *AccountTestService) testKiroAccountConnection(c *gin.Context, account *Account, modelID string) error {
+	ctx := c.Request.Context()
+
+	testModelID := strings.TrimSpace(modelID)
+	if testModelID == "" {
+		testModelID = "claude-sonnet-4-6"
+	}
+	if mappedModel := account.GetMappedModel(testModelID); strings.TrimSpace(mappedModel) != "" {
+		testModelID = mappedModel
+	}
+
+	if account.Type != AccountTypeOAuth {
+		return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Kiro account type: %s", account.Type))
+	}
+
+	if s.kiroTokenProvider == nil {
+		return s.sendErrorAndEnd(c, "Kiro token provider not configured")
+	}
+
+	accessToken, err := s.kiroTokenProvider.GetAccessToken(ctx, account)
+	if err != nil {
+		return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to get Kiro access token: %s", err.Error()))
+	}
+
+	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	c.Writer.Header().Set("Cache-Control", "no-cache")
+	c.Writer.Header().Set("Connection", "keep-alive")
+	c.Writer.Header().Set("X-Accel-Buffering", "no")
+	c.Writer.Flush()
+
+	payload, err := createTestPayload(testModelID)
+	if err != nil {
+		return s.sendErrorAndEnd(c, "Failed to create test payload")
+	}
+	payloadBytes, _ := json.Marshal(payload)
+
+	s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
+
+	resp, err := s.executeKiroTestUpstream(ctx, account, payloadBytes, testModelID, accessToken)
+	if err != nil {
+		return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		defer func() { _ = resp.Body.Close() }()
+		body, _ := io.ReadAll(resp.Body)
+		return s.sendErrorAndEnd(c, formatKiroTestError(resp.StatusCode, body, testModelID, account))
+	}
+
+	pr, pw := io.Pipe()
+	go func() {
+		defer func() { _ = resp.Body.Close() }()
+		_, streamErr := kiropkg.StreamEventStreamAsAnthropic(ctx, resp.Body, pw, testModelID, estimateKiroInputTokens(payloadBytes))
+		if streamErr != nil {
+			_ = pw.CloseWithError(streamErr)
+			return
+		}
+		_ = pw.Close()
+	}()
+
+	return s.processClaudeStream(c, pr)
+}
+
+func formatKiroTestError(statusCode int, body []byte, requestedModel string, account *Account) string {
+	return fmt.Sprintf("API returned %d: %s", statusCode, string(body))
+}
+
+func (s *AccountTestService) executeKiroTestUpstream(ctx context.Context, account *Account, anthropicBody []byte, mappedModel, token string) (*http.Response, error) {
+	modelID := kiropkg.MapModel(mappedModel)
+	currentToken := token
+	buildResult, err := buildKiroPayloadForAccountWithRepo(ctx, s.accountRepo, account, anthropicBody, modelID, currentToken, mappedModel, nil)
+	if err != nil {
+		return nil, err
+	}
+	payload := buildResult.Payload
+
+	endpoints := buildKiroEndpoints(account)
+	proxyURL := kiroProxyURL(account)
+	tlsProfile := s.tlsFPProfileService.ResolveTLSProfile(account)
+	accountKey := buildKiroAccountKey(account)
+	maxRetries := 2
+	for idx, endpoint := range endpoints {
+		for attempt := 0; attempt <= maxRetries; attempt++ {
+			req, err := newKiroJSONRequest(ctx, endpoint.URL, payload, currentToken, accountKey, buildKiroMachineID(account), endpoint.AmzTarget, account)
+			if err != nil {
+				return nil, err
+			}
+
+			resp, err := s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, tlsProfile)
+			if err != nil {
+				return nil, err
+			}
+
+			if resp.StatusCode == http.StatusTooManyRequests || (resp.StatusCode >= 500 && resp.StatusCode < 600) {
+				if idx+1 < len(endpoints) {
+					_ = resp.Body.Close()
+					break
+				}
+				return resp, nil
+			}
+
+			if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
+				respBody, readErr := io.ReadAll(resp.Body)
+				_ = resp.Body.Close()
+				if readErr != nil {
+					return nil, readErr
+				}
+
+				if s.kiroTokenProvider != nil && (resp.StatusCode == http.StatusUnauthorized || isKiroTokenErrorBody(respBody)) && attempt < maxRetries {
+					refreshedToken, refreshErr := s.kiroTokenProvider.ForceRefreshAccessToken(ctx, account)
+					if refreshErr == nil && strings.TrimSpace(refreshedToken) != "" {
+						currentToken = refreshedToken
+						accountKey = buildKiroAccountKey(account)
+						buildResult, err = buildKiroPayloadForAccountWithRepo(ctx, s.accountRepo, account, anthropicBody, modelID, currentToken, mappedModel, nil)
+						if err != nil {
+							return nil, err
+						}
+						payload = buildResult.Payload
+						continue
+					}
+				}
+
+				resetHTTPResponseBody(resp, respBody)
+				return resp, nil
+			}
+
+			if resp.StatusCode == http.StatusBadRequest {
+				respBody, readErr := io.ReadAll(resp.Body)
+				_ = resp.Body.Close()
+				if readErr != nil {
+					return nil, readErr
+				}
+				resetHTTPResponseBody(resp, respBody)
+				return resp, nil
+			}
+
+			return resp, nil
+		}
+	}
+
+	return nil, fmt.Errorf("kiro upstream endpoints exhausted")
+}
+
 // testBedrockAccountConnection tests a Bedrock (SigV4 or API Key) account using non-streaming invoke
 func (s *AccountTestService) testBedrockAccountConnection(c *gin.Context, ctx context.Context, account *Account, testModelID string) error {
 	region := bedrockRuntimeRegion(account)