feat(kiro): 新增 AWS Builder ID 授权码流程认证及用户邮箱获取增强

Amp-Thread-ID: https://ampcode.com/threads/T-019b2ecc-fb2d-713f-b30d-1196c7dce3e2 Co-authored-by: Amp <amp@ampcode.com>
2026-07-14 10:35:24 +00:00 · 2025-12-18 08:16:36 +08:00
parent 54acd69e9d
commit cf9a246d53
6 changed files with 695 additions and 7 deletions
@@ -0,0 +1,166 @@
+// Package kiro provides CodeWhisperer API client for fetching user info.
+package kiro
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	codeWhispererAPI = "https://codewhisperer.us-east-1.amazonaws.com"
+	kiroVersion      = "0.6.18"
+)
+
+// CodeWhispererClient handles CodeWhisperer API calls.
+type CodeWhispererClient struct {
+	httpClient *http.Client
+	machineID  string
+}
+
+// UsageLimitsResponse represents the getUsageLimits API response.
+type UsageLimitsResponse struct {
+	DaysUntilReset     *int                `json:"daysUntilReset,omitempty"`
+	NextDateReset      *float64            `json:"nextDateReset,omitempty"`
+	UserInfo           *UserInfo           `json:"userInfo,omitempty"`
+	SubscriptionInfo   *SubscriptionInfo   `json:"subscriptionInfo,omitempty"`
+	UsageBreakdownList []UsageBreakdown    `json:"usageBreakdownList,omitempty"`
+}
+
+// UserInfo contains user information from the API.
+type UserInfo struct {
+	Email  string `json:"email,omitempty"`
+	UserID string `json:"userId,omitempty"`
+}
+
+// SubscriptionInfo contains subscription details.
+type SubscriptionInfo struct {
+	SubscriptionTitle string `json:"subscriptionTitle,omitempty"`
+	Type              string `json:"type,omitempty"`
+}
+
+// UsageBreakdown contains usage details.
+type UsageBreakdown struct {
+	UsageLimit                 *int     `json:"usageLimit,omitempty"`
+	CurrentUsage               *int     `json:"currentUsage,omitempty"`
+	UsageLimitWithPrecision    *float64 `json:"usageLimitWithPrecision,omitempty"`
+	CurrentUsageWithPrecision  *float64 `json:"currentUsageWithPrecision,omitempty"`
+	NextDateReset              *float64 `json:"nextDateReset,omitempty"`
+	DisplayName                string   `json:"displayName,omitempty"`
+	ResourceType               string   `json:"resourceType,omitempty"`
+}
+
+// NewCodeWhispererClient creates a new CodeWhisperer client.
+func NewCodeWhispererClient(cfg *config.Config, machineID string) *CodeWhispererClient {
+	client := &http.Client{Timeout: 30 * time.Second}
+	if cfg != nil {
+		client = util.SetProxy(&cfg.SDKConfig, client)
+	}
+	if machineID == "" {
+		machineID = uuid.New().String()
+	}
+	return &CodeWhispererClient{
+		httpClient: client,
+		machineID:  machineID,
+	}
+}
+
+// generateInvocationID generates a unique invocation ID.
+func generateInvocationID() string {
+	return uuid.New().String()
+}
+
+// GetUsageLimits fetches usage limits and user info from CodeWhisperer API.
+// This is the recommended way to get user email after login.
+func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken string) (*UsageLimitsResponse, error) {
+	url := fmt.Sprintf("%s/getUsageLimits?isEmailRequired=true&origin=AI_EDITOR&resourceType=AGENTIC_REQUEST", codeWhispererAPI)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	// Set headers to match Kiro IDE
+	xAmzUserAgent := fmt.Sprintf("aws-sdk-js/1.0.0 KiroIDE-%s-%s", kiroVersion, c.machineID)
+	userAgent := fmt.Sprintf("aws-sdk-js/1.0.0 ua/2.1 os/windows lang/js md/nodejs#20.16.0 api/codewhispererruntime#1.0.0 m/E KiroIDE-%s-%s", kiroVersion, c.machineID)
+
+	req.Header.Set("Authorization", "Bearer "+accessToken)
+	req.Header.Set("x-amz-user-agent", xAmzUserAgent)
+	req.Header.Set("User-Agent", userAgent)
+	req.Header.Set("amz-sdk-invocation-id", generateInvocationID())
+	req.Header.Set("amz-sdk-request", "attempt=1; max=1")
+	req.Header.Set("Connection", "close")
+
+	log.Debugf("codewhisperer: GET %s", url)
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response: %w", err)
+	}
+
+	log.Debugf("codewhisperer: status=%d, body=%s", resp.StatusCode, string(body))
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(body))
+	}
+
+	var result UsageLimitsResponse
+	if err := json.Unmarshal(body, &result); err != nil {
+		return nil, fmt.Errorf("failed to parse response: %w", err)
+	}
+
+	return &result, nil
+}
+
+// FetchUserEmailFromAPI fetches user email using CodeWhisperer getUsageLimits API.
+// This is more reliable than JWT parsing as it uses the official API.
+func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessToken string) string {
+	resp, err := c.GetUsageLimits(ctx, accessToken)
+	if err != nil {
+		log.Debugf("codewhisperer: failed to get usage limits: %v", err)
+		return ""
+	}
+
+	if resp.UserInfo != nil && resp.UserInfo.Email != "" {
+		log.Debugf("codewhisperer: got email from API: %s", resp.UserInfo.Email)
+		return resp.UserInfo.Email
+	}
+
+	log.Debugf("codewhisperer: no email in response")
+	return ""
+}
+
+// FetchUserEmailWithFallback fetches user email with multiple fallback methods.
+// Priority: 1. CodeWhisperer API  2. userinfo endpoint  3. JWT parsing
+func FetchUserEmailWithFallback(ctx context.Context, cfg *config.Config, accessToken string) string {
+	// Method 1: Try CodeWhisperer API (most reliable)
+	cwClient := NewCodeWhispererClient(cfg, "")
+	email := cwClient.FetchUserEmailFromAPI(ctx, accessToken)
+	if email != "" {
+		return email
+	}
+
+	// Method 2: Try SSO OIDC userinfo endpoint
+	ssoClient := NewSSOOIDCClient(cfg)
+	email = ssoClient.FetchUserEmail(ctx, accessToken)
+	if email != "" {
+		return email
+	}
+
+	// Method 3: Fallback to JWT parsing
+	return ExtractEmailFromJWT(accessToken)
+}