From 0e77e93e5dd7fc983fb7bbf88b000104bdfab23f Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Sun, 18 Jan 2026 15:04:29 +0800 Subject: [PATCH 01/10] feat: add Kiro OAuth web, rate limiter, metrics, fingerprint, background refresh and model converter --- internal/api/server.go | 6 + internal/auth/kiro/aws.go | 83 + internal/auth/kiro/aws.go.bak | 305 ++++ internal/auth/kiro/background_refresh.go | 192 +++ internal/auth/kiro/cooldown.go | 112 ++ internal/auth/kiro/cooldown_test.go | 240 +++ internal/auth/kiro/fingerprint.go | 197 +++ internal/auth/kiro/fingerprint_test.go | 227 +++ internal/auth/kiro/jitter.go | 174 +++ internal/auth/kiro/metrics.go | 187 +++ internal/auth/kiro/metrics_test.go | 301 ++++ internal/auth/kiro/oauth.go | 2 + internal/auth/kiro/oauth_web.go | 825 ++++++++++ internal/auth/kiro/oauth_web.go.bak | 385 +++++ internal/auth/kiro/oauth_web_templates.go | 732 +++++++++ internal/auth/kiro/rate_limiter.go | 316 ++++ internal/auth/kiro/rate_limiter_singleton.go | 46 + internal/auth/kiro/rate_limiter_test.go | 304 ++++ internal/auth/kiro/social_auth.go | 270 ++-- internal/auth/kiro/sso_oidc.go | 19 +- internal/auth/kiro/sso_oidc.go.bak | 1371 +++++++++++++++++ internal/auth/kiro/token.go | 17 + internal/auth/kiro/usage_checker.go | 243 +++ internal/registry/kiro_model_converter.go | 303 ++++ internal/runtime/executor/kiro_executor.go | 503 +++++- .../translator/kiro/common/utf8_stream.go | 97 ++ .../kiro/common/utf8_stream_test.go | 402 +++++ internal/watcher/events.go | 4 +- sdk/auth/kiro.go | 144 +- sdk/auth/kiro.go.bak | 470 ++++++ sdk/cliproxy/service.go | 201 ++- test_api.py | 452 ++++++ test_auth_diff.go | 273 ++++ test_auth_idc_go1.go | 323 ++++ test_auth_js_style.go | 237 +++ test_kiro_debug.go | 348 +++++ test_proxy_debug.go | 367 +++++ 37 files changed, 10396 insertions(+), 282 deletions(-) create mode 100644 internal/auth/kiro/aws.go.bak create mode 100644 internal/auth/kiro/background_refresh.go create mode 100644 internal/auth/kiro/cooldown.go create mode 100644 internal/auth/kiro/cooldown_test.go create mode 100644 internal/auth/kiro/fingerprint.go create mode 100644 internal/auth/kiro/fingerprint_test.go create mode 100644 internal/auth/kiro/jitter.go create mode 100644 internal/auth/kiro/metrics.go create mode 100644 internal/auth/kiro/metrics_test.go create mode 100644 internal/auth/kiro/oauth_web.go create mode 100644 internal/auth/kiro/oauth_web.go.bak create mode 100644 internal/auth/kiro/oauth_web_templates.go create mode 100644 internal/auth/kiro/rate_limiter.go create mode 100644 internal/auth/kiro/rate_limiter_singleton.go create mode 100644 internal/auth/kiro/rate_limiter_test.go create mode 100644 internal/auth/kiro/sso_oidc.go.bak create mode 100644 internal/auth/kiro/usage_checker.go create mode 100644 internal/registry/kiro_model_converter.go create mode 100644 internal/translator/kiro/common/utf8_stream.go create mode 100644 internal/translator/kiro/common/utf8_stream_test.go create mode 100644 sdk/auth/kiro.go.bak create mode 100644 test_api.py create mode 100644 test_auth_diff.go create mode 100644 test_auth_idc_go1.go create mode 100644 test_auth_js_style.go create mode 100644 test_kiro_debug.go create mode 100644 test_proxy_debug.go diff --git a/internal/api/server.go b/internal/api/server.go index 4df42ec8..40c41a94 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -23,6 +23,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v6/internal/api/middleware" "github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules" ampmodule "github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules/amp" + "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/logging" "github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset" @@ -295,6 +296,11 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk s.registerManagementRoutes() } + // === CLIProxyAPIPlus 扩展: 注册 Kiro OAuth Web 路由 === + kiroOAuthHandler := kiro.NewOAuthWebHandler(cfg) + kiroOAuthHandler.RegisterRoutes(engine) + log.Info("Kiro OAuth Web routes registered at /v0/oauth/kiro/*") + if optionState.keepAliveEnabled { s.enableKeepAlive(optionState.keepAliveTimeout, optionState.keepAliveOnTimeout) } diff --git a/internal/auth/kiro/aws.go b/internal/auth/kiro/aws.go index ba73af4d..d266b9bf 100644 --- a/internal/auth/kiro/aws.go +++ b/internal/auth/kiro/aws.go @@ -5,10 +5,12 @@ package kiro import ( "encoding/base64" "encoding/json" + "errors" "fmt" "os" "path/filepath" "strings" + "time" ) // PKCECodes holds PKCE verification codes for OAuth2 PKCE flow @@ -85,6 +87,87 @@ type KiroModel struct { // KiroIDETokenFile is the default path to Kiro IDE's token file const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json" +// Default retry configuration for file reading +const ( + defaultTokenReadMaxAttempts = 10 // Maximum retry attempts + defaultTokenReadBaseDelay = 50 * time.Millisecond // Base delay between retries +) + +// isTransientFileError checks if the error is a transient file access error +// that may be resolved by retrying (e.g., file locked by another process on Windows). +func isTransientFileError(err error) bool { + if err == nil { + return false + } + + // Check for OS-level file access errors (Windows sharing violation, etc.) + var pathErr *os.PathError + if errors.As(err, &pathErr) { + // Windows sharing violation (ERROR_SHARING_VIOLATION = 32) + // Windows lock violation (ERROR_LOCK_VIOLATION = 33) + errStr := pathErr.Err.Error() + if strings.Contains(errStr, "being used by another process") || + strings.Contains(errStr, "sharing violation") || + strings.Contains(errStr, "lock violation") { + return true + } + } + + // Check error message for common transient patterns + errMsg := strings.ToLower(err.Error()) + transientPatterns := []string{ + "being used by another process", + "sharing violation", + "lock violation", + "access is denied", + "unexpected end of json", + "unexpected eof", + } + for _, pattern := range transientPatterns { + if strings.Contains(errMsg, pattern) { + return true + } + } + + return false +} + +// LoadKiroIDETokenWithRetry loads token data from Kiro IDE's token file with retry logic. +// This handles transient file access errors (e.g., file locked by Kiro IDE during write). +// maxAttempts: maximum number of retry attempts (default 10 if <= 0) +// baseDelay: base delay between retries with exponential backoff (default 50ms if <= 0) +func LoadKiroIDETokenWithRetry(maxAttempts int, baseDelay time.Duration) (*KiroTokenData, error) { + if maxAttempts <= 0 { + maxAttempts = defaultTokenReadMaxAttempts + } + if baseDelay <= 0 { + baseDelay = defaultTokenReadBaseDelay + } + + var lastErr error + for attempt := 0; attempt < maxAttempts; attempt++ { + token, err := LoadKiroIDEToken() + if err == nil { + return token, nil + } + lastErr = err + + // Only retry for transient errors + if !isTransientFileError(err) { + return nil, err + } + + // Exponential backoff: delay * 2^attempt, capped at 500ms + delay := baseDelay * time.Duration(1< 500*time.Millisecond { + delay = 500 * time.Millisecond + } + time.Sleep(delay) + } + + return nil, fmt.Errorf("failed to read token file after %d attempts: %w", maxAttempts, lastErr) +} + // LoadKiroIDEToken loads token data from Kiro IDE's token file. func LoadKiroIDEToken() (*KiroTokenData, error) { homeDir, err := os.UserHomeDir() diff --git a/internal/auth/kiro/aws.go.bak b/internal/auth/kiro/aws.go.bak new file mode 100644 index 00000000..ba73af4d --- /dev/null +++ b/internal/auth/kiro/aws.go.bak @@ -0,0 +1,305 @@ +// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API. +// It includes interfaces and implementations for token storage and authentication methods. +package kiro + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" +) + +// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow +type PKCECodes struct { + // CodeVerifier is the cryptographically random string used to correlate + // the authorization request to the token request + CodeVerifier string `json:"code_verifier"` + // CodeChallenge is the SHA256 hash of the code verifier, base64url-encoded + CodeChallenge string `json:"code_challenge"` +} + +// KiroTokenData holds OAuth token information from AWS CodeWhisperer (Kiro) +type KiroTokenData struct { + // AccessToken is the OAuth2 access token for API access + AccessToken string `json:"accessToken"` + // RefreshToken is used to obtain new access tokens + RefreshToken string `json:"refreshToken"` + // ProfileArn is the AWS CodeWhisperer profile ARN + ProfileArn string `json:"profileArn"` + // ExpiresAt is the timestamp when the token expires + ExpiresAt string `json:"expiresAt"` + // AuthMethod indicates the authentication method used (e.g., "builder-id", "social") + AuthMethod string `json:"authMethod"` + // Provider indicates the OAuth provider (e.g., "AWS", "Google") + Provider string `json:"provider"` + // ClientID is the OIDC client ID (needed for token refresh) + ClientID string `json:"clientId,omitempty"` + // ClientSecret is the OIDC client secret (needed for token refresh) + ClientSecret string `json:"clientSecret,omitempty"` + // Email is the user's email address (used for file naming) + Email string `json:"email,omitempty"` + // StartURL is the IDC/Identity Center start URL (only for IDC auth method) + StartURL string `json:"startUrl,omitempty"` + // Region is the AWS region for IDC authentication (only for IDC auth method) + Region string `json:"region,omitempty"` +} + +// KiroAuthBundle aggregates authentication data after OAuth flow completion +type KiroAuthBundle struct { + // TokenData contains the OAuth tokens from the authentication flow + TokenData KiroTokenData `json:"token_data"` + // LastRefresh is the timestamp of the last token refresh + LastRefresh string `json:"last_refresh"` +} + +// KiroUsageInfo represents usage information from CodeWhisperer API +type KiroUsageInfo struct { + // SubscriptionTitle is the subscription plan name (e.g., "KIRO FREE") + SubscriptionTitle string `json:"subscription_title"` + // CurrentUsage is the current credit usage + CurrentUsage float64 `json:"current_usage"` + // UsageLimit is the maximum credit limit + UsageLimit float64 `json:"usage_limit"` + // NextReset is the timestamp of the next usage reset + NextReset string `json:"next_reset"` +} + +// KiroModel represents a model available through the CodeWhisperer API +type KiroModel struct { + // ModelID is the unique identifier for the model + ModelID string `json:"modelId"` + // ModelName is the human-readable name + ModelName string `json:"modelName"` + // Description is the model description + Description string `json:"description"` + // RateMultiplier is the credit multiplier for this model + RateMultiplier float64 `json:"rateMultiplier"` + // RateUnit is the unit for rate calculation (e.g., "credit") + RateUnit string `json:"rateUnit"` + // MaxInputTokens is the maximum input token limit + MaxInputTokens int `json:"maxInputTokens,omitempty"` +} + +// KiroIDETokenFile is the default path to Kiro IDE's token file +const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json" + +// LoadKiroIDEToken loads token data from Kiro IDE's token file. +func LoadKiroIDEToken() (*KiroTokenData, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("failed to get home directory: %w", err) + } + + tokenPath := filepath.Join(homeDir, KiroIDETokenFile) + data, err := os.ReadFile(tokenPath) + if err != nil { + return nil, fmt.Errorf("failed to read Kiro IDE token file (%s): %w", tokenPath, err) + } + + var token KiroTokenData + if err := json.Unmarshal(data, &token); err != nil { + return nil, fmt.Errorf("failed to parse Kiro IDE token: %w", err) + } + + if token.AccessToken == "" { + return nil, fmt.Errorf("access token is empty in Kiro IDE token file") + } + + return &token, nil +} + +// LoadKiroTokenFromPath loads token data from a custom path. +// This supports multiple accounts by allowing different token files. +func LoadKiroTokenFromPath(tokenPath string) (*KiroTokenData, error) { + // Expand ~ to home directory + if len(tokenPath) > 0 && tokenPath[0] == '~' { + homeDir, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("failed to get home directory: %w", err) + } + tokenPath = filepath.Join(homeDir, tokenPath[1:]) + } + + data, err := os.ReadFile(tokenPath) + if err != nil { + return nil, fmt.Errorf("failed to read token file (%s): %w", tokenPath, err) + } + + var token KiroTokenData + if err := json.Unmarshal(data, &token); err != nil { + return nil, fmt.Errorf("failed to parse token file: %w", err) + } + + if token.AccessToken == "" { + return nil, fmt.Errorf("access token is empty in token file") + } + + return &token, nil +} + +// ListKiroTokenFiles lists all Kiro token files in the cache directory. +// This supports multiple accounts by finding all token files. +func ListKiroTokenFiles() ([]string, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("failed to get home directory: %w", err) + } + + cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache") + + // Check if directory exists + if _, err := os.Stat(cacheDir); os.IsNotExist(err) { + return nil, nil // No token files + } + + entries, err := os.ReadDir(cacheDir) + if err != nil { + return nil, fmt.Errorf("failed to read cache directory: %w", err) + } + + var tokenFiles []string + for _, entry := range entries { + if entry.IsDir() { + continue + } + name := entry.Name() + // Look for kiro token files only (avoid matching unrelated AWS SSO cache files) + if strings.HasSuffix(name, ".json") && strings.HasPrefix(name, "kiro") { + tokenFiles = append(tokenFiles, filepath.Join(cacheDir, name)) + } + } + + return tokenFiles, nil +} + +// LoadAllKiroTokens loads all Kiro tokens from the cache directory. +// This supports multiple accounts. +func LoadAllKiroTokens() ([]*KiroTokenData, error) { + files, err := ListKiroTokenFiles() + if err != nil { + return nil, err + } + + var tokens []*KiroTokenData + for _, file := range files { + token, err := LoadKiroTokenFromPath(file) + if err != nil { + // Skip invalid token files + continue + } + tokens = append(tokens, token) + } + + return tokens, nil +} + +// JWTClaims represents the claims we care about from a JWT token. +// JWT tokens from Kiro/AWS contain user information in the payload. +type JWTClaims struct { + Email string `json:"email,omitempty"` + Sub string `json:"sub,omitempty"` + PreferredUser string `json:"preferred_username,omitempty"` + Name string `json:"name,omitempty"` + Iss string `json:"iss,omitempty"` +} + +// ExtractEmailFromJWT extracts the user's email from a JWT access token. +// JWT tokens typically have format: header.payload.signature +// The payload is base64url-encoded JSON containing user claims. +func ExtractEmailFromJWT(accessToken string) string { + if accessToken == "" { + return "" + } + + // JWT format: header.payload.signature + parts := strings.Split(accessToken, ".") + if len(parts) != 3 { + return "" + } + + // Decode the payload (second part) + payload := parts[1] + + // Add padding if needed (base64url requires padding) + switch len(payload) % 4 { + case 2: + payload += "==" + case 3: + payload += "=" + } + + decoded, err := base64.URLEncoding.DecodeString(payload) + if err != nil { + // Try RawURLEncoding (no padding) + decoded, err = base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return "" + } + } + + var claims JWTClaims + if err := json.Unmarshal(decoded, &claims); err != nil { + return "" + } + + // Return email if available + if claims.Email != "" { + return claims.Email + } + + // Fallback to preferred_username (some providers use this) + if claims.PreferredUser != "" && strings.Contains(claims.PreferredUser, "@") { + return claims.PreferredUser + } + + // Fallback to sub if it looks like an email + if claims.Sub != "" && strings.Contains(claims.Sub, "@") { + return claims.Sub + } + + return "" +} + +// SanitizeEmailForFilename sanitizes an email address for use in a filename. +// Replaces special characters with underscores and prevents path traversal attacks. +// Also handles URL-encoded characters to prevent encoded path traversal attempts. +func SanitizeEmailForFilename(email string) string { + if email == "" { + return "" + } + + result := email + + // First, handle URL-encoded path traversal attempts (%2F, %2E, %5C, etc.) + // This prevents encoded characters from bypassing the sanitization. + // Note: We replace % last to catch any remaining encodings including double-encoding (%252F) + result = strings.ReplaceAll(result, "%2F", "_") // / + result = strings.ReplaceAll(result, "%2f", "_") + result = strings.ReplaceAll(result, "%5C", "_") // \ + result = strings.ReplaceAll(result, "%5c", "_") + result = strings.ReplaceAll(result, "%2E", "_") // . + result = strings.ReplaceAll(result, "%2e", "_") + result = strings.ReplaceAll(result, "%00", "_") // null byte + result = strings.ReplaceAll(result, "%", "_") // Catch remaining % to prevent double-encoding attacks + + // Replace characters that are problematic in filenames + // Keep @ and . in middle but replace other special characters + for _, char := range []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " ", "\x00"} { + result = strings.ReplaceAll(result, char, "_") + } + + // Prevent path traversal: replace leading dots in each path component + // This handles cases like "../../../etc/passwd" → "_.._.._.._etc_passwd" + parts := strings.Split(result, "_") + for i, part := range parts { + for strings.HasPrefix(part, ".") { + part = "_" + part[1:] + } + parts[i] = part + } + result = strings.Join(parts, "_") + + return result +} diff --git a/internal/auth/kiro/background_refresh.go b/internal/auth/kiro/background_refresh.go new file mode 100644 index 00000000..3fecc417 --- /dev/null +++ b/internal/auth/kiro/background_refresh.go @@ -0,0 +1,192 @@ +package kiro + +import ( + "context" + "log" + "sync" + "time" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + "golang.org/x/sync/semaphore" +) + +type Token struct { + ID string + AccessToken string + RefreshToken string + ExpiresAt time.Time + LastVerified time.Time + ClientID string + ClientSecret string + AuthMethod string + Provider string + StartURL string + Region string +} + +type TokenRepository interface { + FindOldestUnverified(limit int) []*Token + UpdateToken(token *Token) error +} + +type RefresherOption func(*BackgroundRefresher) + +func WithInterval(interval time.Duration) RefresherOption { + return func(r *BackgroundRefresher) { + r.interval = interval + } +} + +func WithBatchSize(size int) RefresherOption { + return func(r *BackgroundRefresher) { + r.batchSize = size + } +} + +func WithConcurrency(concurrency int) RefresherOption { + return func(r *BackgroundRefresher) { + r.concurrency = concurrency + } +} + +type BackgroundRefresher struct { + interval time.Duration + batchSize int + concurrency int + tokenRepo TokenRepository + stopCh chan struct{} + wg sync.WaitGroup + oauth *KiroOAuth + ssoClient *SSOOIDCClient +} + +func NewBackgroundRefresher(repo TokenRepository, opts ...RefresherOption) *BackgroundRefresher { + r := &BackgroundRefresher{ + interval: time.Minute, + batchSize: 50, + concurrency: 10, + tokenRepo: repo, + stopCh: make(chan struct{}), + oauth: nil, // Lazy init - will be set when config available + ssoClient: nil, // Lazy init - will be set when config available + } + for _, opt := range opts { + opt(r) + } + return r +} + +// WithConfig sets the configuration for OAuth and SSO clients. +func WithConfig(cfg *config.Config) RefresherOption { + return func(r *BackgroundRefresher) { + r.oauth = NewKiroOAuth(cfg) + r.ssoClient = NewSSOOIDCClient(cfg) + } +} + +func (r *BackgroundRefresher) Start(ctx context.Context) { + r.wg.Add(1) + go func() { + defer r.wg.Done() + ticker := time.NewTicker(r.interval) + defer ticker.Stop() + + r.refreshBatch(ctx) + + for { + select { + case <-ctx.Done(): + return + case <-r.stopCh: + return + case <-ticker.C: + r.refreshBatch(ctx) + } + } + }() +} + +func (r *BackgroundRefresher) Stop() { + close(r.stopCh) + r.wg.Wait() +} + +func (r *BackgroundRefresher) refreshBatch(ctx context.Context) { + tokens := r.tokenRepo.FindOldestUnverified(r.batchSize) + if len(tokens) == 0 { + return + } + + sem := semaphore.NewWeighted(int64(r.concurrency)) + var wg sync.WaitGroup + + for i, token := range tokens { + if i > 0 { + select { + case <-ctx.Done(): + return + case <-r.stopCh: + return + case <-time.After(100 * time.Millisecond): + } + } + + if err := sem.Acquire(ctx, 1); err != nil { + return + } + + wg.Add(1) + go func(t *Token) { + defer wg.Done() + defer sem.Release(1) + r.refreshSingle(ctx, t) + }(token) + } + + wg.Wait() +} + +func (r *BackgroundRefresher) refreshSingle(ctx context.Context, token *Token) { + var newTokenData *KiroTokenData + var err error + + switch token.AuthMethod { + case "idc": + newTokenData, err = r.ssoClient.RefreshTokenWithRegion( + ctx, + token.ClientID, + token.ClientSecret, + token.RefreshToken, + token.Region, + token.StartURL, + ) + case "builder-id": + newTokenData, err = r.ssoClient.RefreshToken( + ctx, + token.ClientID, + token.ClientSecret, + token.RefreshToken, + ) + default: + newTokenData, err = r.oauth.RefreshToken(ctx, token.RefreshToken) + } + + if err != nil { + log.Printf("failed to refresh token %s: %v", token.ID, err) + return + } + + token.AccessToken = newTokenData.AccessToken + token.RefreshToken = newTokenData.RefreshToken + token.LastVerified = time.Now() + + if newTokenData.ExpiresAt != "" { + if expTime, parseErr := time.Parse(time.RFC3339, newTokenData.ExpiresAt); parseErr == nil { + token.ExpiresAt = expTime + } + } + + if err := r.tokenRepo.UpdateToken(token); err != nil { + log.Printf("failed to update token %s: %v", token.ID, err) + } +} diff --git a/internal/auth/kiro/cooldown.go b/internal/auth/kiro/cooldown.go new file mode 100644 index 00000000..c1aabbcb --- /dev/null +++ b/internal/auth/kiro/cooldown.go @@ -0,0 +1,112 @@ +package kiro + +import ( + "sync" + "time" +) + +const ( + CooldownReason429 = "rate_limit_exceeded" + CooldownReasonSuspended = "account_suspended" + CooldownReasonQuotaExhausted = "quota_exhausted" + + DefaultShortCooldown = 1 * time.Minute + MaxShortCooldown = 5 * time.Minute + LongCooldown = 24 * time.Hour +) + +type CooldownManager struct { + mu sync.RWMutex + cooldowns map[string]time.Time + reasons map[string]string +} + +func NewCooldownManager() *CooldownManager { + return &CooldownManager{ + cooldowns: make(map[string]time.Time), + reasons: make(map[string]string), + } +} + +func (cm *CooldownManager) SetCooldown(tokenKey string, duration time.Duration, reason string) { + cm.mu.Lock() + defer cm.mu.Unlock() + cm.cooldowns[tokenKey] = time.Now().Add(duration) + cm.reasons[tokenKey] = reason +} + +func (cm *CooldownManager) IsInCooldown(tokenKey string) bool { + cm.mu.RLock() + defer cm.mu.RUnlock() + endTime, exists := cm.cooldowns[tokenKey] + if !exists { + return false + } + return time.Now().Before(endTime) +} + +func (cm *CooldownManager) GetRemainingCooldown(tokenKey string) time.Duration { + cm.mu.RLock() + defer cm.mu.RUnlock() + endTime, exists := cm.cooldowns[tokenKey] + if !exists { + return 0 + } + remaining := time.Until(endTime) + if remaining < 0 { + return 0 + } + return remaining +} + +func (cm *CooldownManager) GetCooldownReason(tokenKey string) string { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.reasons[tokenKey] +} + +func (cm *CooldownManager) ClearCooldown(tokenKey string) { + cm.mu.Lock() + defer cm.mu.Unlock() + delete(cm.cooldowns, tokenKey) + delete(cm.reasons, tokenKey) +} + +func (cm *CooldownManager) CleanupExpired() { + cm.mu.Lock() + defer cm.mu.Unlock() + now := time.Now() + for tokenKey, endTime := range cm.cooldowns { + if now.After(endTime) { + delete(cm.cooldowns, tokenKey) + delete(cm.reasons, tokenKey) + } + } +} + +func (cm *CooldownManager) StartCleanupRoutine(interval time.Duration, stopCh <-chan struct{}) { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + cm.CleanupExpired() + case <-stopCh: + return + } + } +} + +func CalculateCooldownFor429(retryCount int) time.Duration { + duration := DefaultShortCooldown * time.Duration(1< MaxShortCooldown { + return MaxShortCooldown + } + return duration +} + +func CalculateCooldownUntilNextDay() time.Duration { + now := time.Now() + nextDay := time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()) + return time.Until(nextDay) +} diff --git a/internal/auth/kiro/cooldown_test.go b/internal/auth/kiro/cooldown_test.go new file mode 100644 index 00000000..e0b35df4 --- /dev/null +++ b/internal/auth/kiro/cooldown_test.go @@ -0,0 +1,240 @@ +package kiro + +import ( + "sync" + "testing" + "time" +) + +func TestNewCooldownManager(t *testing.T) { + cm := NewCooldownManager() + if cm == nil { + t.Fatal("expected non-nil CooldownManager") + } + if cm.cooldowns == nil { + t.Error("expected non-nil cooldowns map") + } + if cm.reasons == nil { + t.Error("expected non-nil reasons map") + } +} + +func TestSetCooldown(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("token1", 1*time.Minute, CooldownReason429) + + if !cm.IsInCooldown("token1") { + t.Error("expected token to be in cooldown") + } + if cm.GetCooldownReason("token1") != CooldownReason429 { + t.Errorf("expected reason %s, got %s", CooldownReason429, cm.GetCooldownReason("token1")) + } +} + +func TestIsInCooldown_NotSet(t *testing.T) { + cm := NewCooldownManager() + if cm.IsInCooldown("nonexistent") { + t.Error("expected non-existent token to not be in cooldown") + } +} + +func TestIsInCooldown_Expired(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("token1", 1*time.Millisecond, CooldownReason429) + + time.Sleep(10 * time.Millisecond) + + if cm.IsInCooldown("token1") { + t.Error("expected expired cooldown to return false") + } +} + +func TestGetRemainingCooldown(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("token1", 1*time.Second, CooldownReason429) + + remaining := cm.GetRemainingCooldown("token1") + if remaining <= 0 || remaining > 1*time.Second { + t.Errorf("expected remaining cooldown between 0 and 1s, got %v", remaining) + } +} + +func TestGetRemainingCooldown_NotSet(t *testing.T) { + cm := NewCooldownManager() + remaining := cm.GetRemainingCooldown("nonexistent") + if remaining != 0 { + t.Errorf("expected 0 remaining for non-existent, got %v", remaining) + } +} + +func TestGetRemainingCooldown_Expired(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("token1", 1*time.Millisecond, CooldownReason429) + + time.Sleep(10 * time.Millisecond) + + remaining := cm.GetRemainingCooldown("token1") + if remaining != 0 { + t.Errorf("expected 0 remaining for expired, got %v", remaining) + } +} + +func TestGetCooldownReason(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("token1", 1*time.Minute, CooldownReasonSuspended) + + reason := cm.GetCooldownReason("token1") + if reason != CooldownReasonSuspended { + t.Errorf("expected reason %s, got %s", CooldownReasonSuspended, reason) + } +} + +func TestGetCooldownReason_NotSet(t *testing.T) { + cm := NewCooldownManager() + reason := cm.GetCooldownReason("nonexistent") + if reason != "" { + t.Errorf("expected empty reason for non-existent, got %s", reason) + } +} + +func TestClearCooldown(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("token1", 1*time.Minute, CooldownReason429) + cm.ClearCooldown("token1") + + if cm.IsInCooldown("token1") { + t.Error("expected cooldown to be cleared") + } + if cm.GetCooldownReason("token1") != "" { + t.Error("expected reason to be cleared") + } +} + +func TestClearCooldown_NonExistent(t *testing.T) { + cm := NewCooldownManager() + cm.ClearCooldown("nonexistent") +} + +func TestCleanupExpired(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("expired1", 1*time.Millisecond, CooldownReason429) + cm.SetCooldown("expired2", 1*time.Millisecond, CooldownReason429) + cm.SetCooldown("active", 1*time.Hour, CooldownReason429) + + time.Sleep(10 * time.Millisecond) + cm.CleanupExpired() + + if cm.GetCooldownReason("expired1") != "" { + t.Error("expected expired1 to be cleaned up") + } + if cm.GetCooldownReason("expired2") != "" { + t.Error("expected expired2 to be cleaned up") + } + if cm.GetCooldownReason("active") != CooldownReason429 { + t.Error("expected active to remain") + } +} + +func TestCalculateCooldownFor429_FirstRetry(t *testing.T) { + duration := CalculateCooldownFor429(0) + if duration != DefaultShortCooldown { + t.Errorf("expected %v for retry 0, got %v", DefaultShortCooldown, duration) + } +} + +func TestCalculateCooldownFor429_Exponential(t *testing.T) { + d1 := CalculateCooldownFor429(1) + d2 := CalculateCooldownFor429(2) + + if d2 <= d1 { + t.Errorf("expected d2 > d1, got d1=%v, d2=%v", d1, d2) + } +} + +func TestCalculateCooldownFor429_MaxCap(t *testing.T) { + duration := CalculateCooldownFor429(10) + if duration > MaxShortCooldown { + t.Errorf("expected max %v, got %v", MaxShortCooldown, duration) + } +} + +func TestCalculateCooldownUntilNextDay(t *testing.T) { + duration := CalculateCooldownUntilNextDay() + if duration <= 0 || duration > 24*time.Hour { + t.Errorf("expected duration between 0 and 24h, got %v", duration) + } +} + +func TestCooldownManager_ConcurrentAccess(t *testing.T) { + cm := NewCooldownManager() + const numGoroutines = 50 + const numOperations = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + tokenKey := "token" + string(rune('a'+id%10)) + for j := 0; j < numOperations; j++ { + switch j % 6 { + case 0: + cm.SetCooldown(tokenKey, time.Duration(j)*time.Millisecond, CooldownReason429) + case 1: + cm.IsInCooldown(tokenKey) + case 2: + cm.GetRemainingCooldown(tokenKey) + case 3: + cm.GetCooldownReason(tokenKey) + case 4: + cm.ClearCooldown(tokenKey) + case 5: + cm.CleanupExpired() + } + } + }(i) + } + + wg.Wait() +} + +func TestCooldownReasonConstants(t *testing.T) { + if CooldownReason429 != "rate_limit_exceeded" { + t.Errorf("unexpected CooldownReason429: %s", CooldownReason429) + } + if CooldownReasonSuspended != "account_suspended" { + t.Errorf("unexpected CooldownReasonSuspended: %s", CooldownReasonSuspended) + } + if CooldownReasonQuotaExhausted != "quota_exhausted" { + t.Errorf("unexpected CooldownReasonQuotaExhausted: %s", CooldownReasonQuotaExhausted) + } +} + +func TestDefaultConstants(t *testing.T) { + if DefaultShortCooldown != 1*time.Minute { + t.Errorf("unexpected DefaultShortCooldown: %v", DefaultShortCooldown) + } + if MaxShortCooldown != 5*time.Minute { + t.Errorf("unexpected MaxShortCooldown: %v", MaxShortCooldown) + } + if LongCooldown != 24*time.Hour { + t.Errorf("unexpected LongCooldown: %v", LongCooldown) + } +} + +func TestSetCooldown_OverwritesPrevious(t *testing.T) { + cm := NewCooldownManager() + cm.SetCooldown("token1", 1*time.Hour, CooldownReason429) + cm.SetCooldown("token1", 1*time.Minute, CooldownReasonSuspended) + + reason := cm.GetCooldownReason("token1") + if reason != CooldownReasonSuspended { + t.Errorf("expected reason to be overwritten to %s, got %s", CooldownReasonSuspended, reason) + } + + remaining := cm.GetRemainingCooldown("token1") + if remaining > 1*time.Minute { + t.Errorf("expected remaining <= 1 minute, got %v", remaining) + } +} diff --git a/internal/auth/kiro/fingerprint.go b/internal/auth/kiro/fingerprint.go new file mode 100644 index 00000000..c35e62b2 --- /dev/null +++ b/internal/auth/kiro/fingerprint.go @@ -0,0 +1,197 @@ +package kiro + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "math/rand" + "net/http" + "sync" + "time" +) + +// Fingerprint 多维度指纹信息 +type Fingerprint struct { + SDKVersion string // 1.0.20-1.0.27 + OSType string // darwin/windows/linux + OSVersion string // 10.0.22621 + NodeVersion string // 18.x/20.x/22.x + KiroVersion string // 0.3.x-0.8.x + KiroHash string // SHA256 + AcceptLanguage string + ScreenResolution string // 1920x1080 + ColorDepth int // 24 + HardwareConcurrency int // CPU 核心数 + TimezoneOffset int +} + +// FingerprintManager 指纹管理器 +type FingerprintManager struct { + mu sync.RWMutex + fingerprints map[string]*Fingerprint // tokenKey -> fingerprint + rng *rand.Rand +} + +var ( + sdkVersions = []string{ + "1.0.20", "1.0.21", "1.0.22", "1.0.23", + "1.0.24", "1.0.25", "1.0.26", "1.0.27", + } + osTypes = []string{"darwin", "windows", "linux"} + osVersions = map[string][]string{ + "darwin": {"14.0", "14.1", "14.2", "14.3", "14.4", "14.5", "15.0", "15.1"}, + "windows": {"10.0.19041", "10.0.19042", "10.0.19043", "10.0.19044", "10.0.22621", "10.0.22631"}, + "linux": {"5.15.0", "6.1.0", "6.2.0", "6.5.0", "6.6.0", "6.8.0"}, + } + nodeVersions = []string{ + "18.17.0", "18.18.0", "18.19.0", "18.20.0", + "20.9.0", "20.10.0", "20.11.0", "20.12.0", "20.13.0", + "22.0.0", "22.1.0", "22.2.0", "22.3.0", + } + kiroVersions = []string{ + "0.3.0", "0.3.1", "0.4.0", "0.4.1", "0.5.0", "0.5.1", + "0.6.0", "0.6.1", "0.7.0", "0.7.1", "0.8.0", "0.8.1", + } + acceptLanguages = []string{ + "en-US,en;q=0.9", + "en-GB,en;q=0.9", + "zh-CN,zh;q=0.9,en;q=0.8", + "zh-TW,zh;q=0.9,en;q=0.8", + "ja-JP,ja;q=0.9,en;q=0.8", + "ko-KR,ko;q=0.9,en;q=0.8", + "de-DE,de;q=0.9,en;q=0.8", + "fr-FR,fr;q=0.9,en;q=0.8", + } + screenResolutions = []string{ + "1920x1080", "2560x1440", "3840x2160", + "1366x768", "1440x900", "1680x1050", + "2560x1600", "3440x1440", + } + colorDepths = []int{24, 32} + hardwareConcurrencies = []int{4, 6, 8, 10, 12, 16, 20, 24, 32} + timezoneOffsets = []int{-480, -420, -360, -300, -240, 0, 60, 120, 480, 540} +) + +// NewFingerprintManager 创建指纹管理器 +func NewFingerprintManager() *FingerprintManager { + return &FingerprintManager{ + fingerprints: make(map[string]*Fingerprint), + rng: rand.New(rand.NewSource(time.Now().UnixNano())), + } +} + +// GetFingerprint 获取或生成 Token 关联的指纹 +func (fm *FingerprintManager) GetFingerprint(tokenKey string) *Fingerprint { + fm.mu.RLock() + if fp, exists := fm.fingerprints[tokenKey]; exists { + fm.mu.RUnlock() + return fp + } + fm.mu.RUnlock() + + fm.mu.Lock() + defer fm.mu.Unlock() + + if fp, exists := fm.fingerprints[tokenKey]; exists { + return fp + } + + fp := fm.generateFingerprint(tokenKey) + fm.fingerprints[tokenKey] = fp + return fp +} + +// generateFingerprint 生成新的指纹 +func (fm *FingerprintManager) generateFingerprint(tokenKey string) *Fingerprint { + osType := fm.randomChoice(osTypes) + osVersion := fm.randomChoice(osVersions[osType]) + kiroVersion := fm.randomChoice(kiroVersions) + + fp := &Fingerprint{ + SDKVersion: fm.randomChoice(sdkVersions), + OSType: osType, + OSVersion: osVersion, + NodeVersion: fm.randomChoice(nodeVersions), + KiroVersion: kiroVersion, + AcceptLanguage: fm.randomChoice(acceptLanguages), + ScreenResolution: fm.randomChoice(screenResolutions), + ColorDepth: fm.randomIntChoice(colorDepths), + HardwareConcurrency: fm.randomIntChoice(hardwareConcurrencies), + TimezoneOffset: fm.randomIntChoice(timezoneOffsets), + } + + fp.KiroHash = fm.generateKiroHash(tokenKey, kiroVersion, osType) + return fp +} + +// generateKiroHash 生成 Kiro Hash +func (fm *FingerprintManager) generateKiroHash(tokenKey, kiroVersion, osType string) string { + data := fmt.Sprintf("%s:%s:%s:%d", tokenKey, kiroVersion, osType, time.Now().UnixNano()) + hash := sha256.Sum256([]byte(data)) + return hex.EncodeToString(hash[:]) +} + +// randomChoice 随机选择字符串 +func (fm *FingerprintManager) randomChoice(choices []string) string { + return choices[fm.rng.Intn(len(choices))] +} + +// randomIntChoice 随机选择整数 +func (fm *FingerprintManager) randomIntChoice(choices []int) int { + return choices[fm.rng.Intn(len(choices))] +} + +// ApplyToRequest 将指纹信息应用到 HTTP 请求头 +func (fp *Fingerprint) ApplyToRequest(req *http.Request) { + req.Header.Set("X-Kiro-SDK-Version", fp.SDKVersion) + req.Header.Set("X-Kiro-OS-Type", fp.OSType) + req.Header.Set("X-Kiro-OS-Version", fp.OSVersion) + req.Header.Set("X-Kiro-Node-Version", fp.NodeVersion) + req.Header.Set("X-Kiro-Version", fp.KiroVersion) + req.Header.Set("X-Kiro-Hash", fp.KiroHash) + req.Header.Set("Accept-Language", fp.AcceptLanguage) + req.Header.Set("X-Screen-Resolution", fp.ScreenResolution) + req.Header.Set("X-Color-Depth", fmt.Sprintf("%d", fp.ColorDepth)) + req.Header.Set("X-Hardware-Concurrency", fmt.Sprintf("%d", fp.HardwareConcurrency)) + req.Header.Set("X-Timezone-Offset", fmt.Sprintf("%d", fp.TimezoneOffset)) +} + +// RemoveFingerprint 移除 Token 关联的指纹 +func (fm *FingerprintManager) RemoveFingerprint(tokenKey string) { + fm.mu.Lock() + defer fm.mu.Unlock() + delete(fm.fingerprints, tokenKey) +} + +// Count 返回当前管理的指纹数量 +func (fm *FingerprintManager) Count() int { + fm.mu.RLock() + defer fm.mu.RUnlock() + return len(fm.fingerprints) +} + +// BuildUserAgent 构建 User-Agent 字符串 (Kiro IDE 风格) +// 格式: aws-sdk-js/{SDKVersion} ua/2.1 os/{OSType}#{OSVersion} lang/js md/nodejs#{NodeVersion} api/codewhispererstreaming#{SDKVersion} m/E KiroIDE-{KiroVersion}-{KiroHash} +func (fp *Fingerprint) BuildUserAgent() string { + return fmt.Sprintf( + "aws-sdk-js/%s ua/2.1 os/%s#%s lang/js md/nodejs#%s api/codewhispererstreaming#%s m/E KiroIDE-%s-%s", + fp.SDKVersion, + fp.OSType, + fp.OSVersion, + fp.NodeVersion, + fp.SDKVersion, + fp.KiroVersion, + fp.KiroHash, + ) +} + +// BuildAmzUserAgent 构建 X-Amz-User-Agent 字符串 +// 格式: aws-sdk-js/{SDKVersion} KiroIDE-{KiroVersion}-{KiroHash} +func (fp *Fingerprint) BuildAmzUserAgent() string { + return fmt.Sprintf( + "aws-sdk-js/%s KiroIDE-%s-%s", + fp.SDKVersion, + fp.KiroVersion, + fp.KiroHash, + ) +} diff --git a/internal/auth/kiro/fingerprint_test.go b/internal/auth/kiro/fingerprint_test.go new file mode 100644 index 00000000..e0ae51f2 --- /dev/null +++ b/internal/auth/kiro/fingerprint_test.go @@ -0,0 +1,227 @@ +package kiro + +import ( + "net/http" + "sync" + "testing" +) + +func TestNewFingerprintManager(t *testing.T) { + fm := NewFingerprintManager() + if fm == nil { + t.Fatal("expected non-nil FingerprintManager") + } + if fm.fingerprints == nil { + t.Error("expected non-nil fingerprints map") + } + if fm.rng == nil { + t.Error("expected non-nil rng") + } +} + +func TestGetFingerprint_NewToken(t *testing.T) { + fm := NewFingerprintManager() + fp := fm.GetFingerprint("token1") + + if fp == nil { + t.Fatal("expected non-nil Fingerprint") + } + if fp.SDKVersion == "" { + t.Error("expected non-empty SDKVersion") + } + if fp.OSType == "" { + t.Error("expected non-empty OSType") + } + if fp.OSVersion == "" { + t.Error("expected non-empty OSVersion") + } + if fp.NodeVersion == "" { + t.Error("expected non-empty NodeVersion") + } + if fp.KiroVersion == "" { + t.Error("expected non-empty KiroVersion") + } + if fp.KiroHash == "" { + t.Error("expected non-empty KiroHash") + } + if fp.AcceptLanguage == "" { + t.Error("expected non-empty AcceptLanguage") + } + if fp.ScreenResolution == "" { + t.Error("expected non-empty ScreenResolution") + } + if fp.ColorDepth == 0 { + t.Error("expected non-zero ColorDepth") + } + if fp.HardwareConcurrency == 0 { + t.Error("expected non-zero HardwareConcurrency") + } +} + +func TestGetFingerprint_SameTokenReturnsSameFingerprint(t *testing.T) { + fm := NewFingerprintManager() + fp1 := fm.GetFingerprint("token1") + fp2 := fm.GetFingerprint("token1") + + if fp1 != fp2 { + t.Error("expected same fingerprint for same token") + } +} + +func TestGetFingerprint_DifferentTokens(t *testing.T) { + fm := NewFingerprintManager() + fp1 := fm.GetFingerprint("token1") + fp2 := fm.GetFingerprint("token2") + + if fp1 == fp2 { + t.Error("expected different fingerprints for different tokens") + } +} + +func TestRemoveFingerprint(t *testing.T) { + fm := NewFingerprintManager() + fm.GetFingerprint("token1") + if fm.Count() != 1 { + t.Fatalf("expected count 1, got %d", fm.Count()) + } + + fm.RemoveFingerprint("token1") + if fm.Count() != 0 { + t.Errorf("expected count 0, got %d", fm.Count()) + } +} + +func TestRemoveFingerprint_NonExistent(t *testing.T) { + fm := NewFingerprintManager() + fm.RemoveFingerprint("nonexistent") + if fm.Count() != 0 { + t.Errorf("expected count 0, got %d", fm.Count()) + } +} + +func TestCount(t *testing.T) { + fm := NewFingerprintManager() + if fm.Count() != 0 { + t.Errorf("expected count 0, got %d", fm.Count()) + } + + fm.GetFingerprint("token1") + fm.GetFingerprint("token2") + fm.GetFingerprint("token3") + + if fm.Count() != 3 { + t.Errorf("expected count 3, got %d", fm.Count()) + } +} + +func TestApplyToRequest(t *testing.T) { + fm := NewFingerprintManager() + fp := fm.GetFingerprint("token1") + + req, _ := http.NewRequest("GET", "http://example.com", nil) + fp.ApplyToRequest(req) + + if req.Header.Get("X-Kiro-SDK-Version") != fp.SDKVersion { + t.Error("X-Kiro-SDK-Version header mismatch") + } + if req.Header.Get("X-Kiro-OS-Type") != fp.OSType { + t.Error("X-Kiro-OS-Type header mismatch") + } + if req.Header.Get("X-Kiro-OS-Version") != fp.OSVersion { + t.Error("X-Kiro-OS-Version header mismatch") + } + if req.Header.Get("X-Kiro-Node-Version") != fp.NodeVersion { + t.Error("X-Kiro-Node-Version header mismatch") + } + if req.Header.Get("X-Kiro-Version") != fp.KiroVersion { + t.Error("X-Kiro-Version header mismatch") + } + if req.Header.Get("X-Kiro-Hash") != fp.KiroHash { + t.Error("X-Kiro-Hash header mismatch") + } + if req.Header.Get("Accept-Language") != fp.AcceptLanguage { + t.Error("Accept-Language header mismatch") + } + if req.Header.Get("X-Screen-Resolution") != fp.ScreenResolution { + t.Error("X-Screen-Resolution header mismatch") + } +} + +func TestGetFingerprint_OSVersionMatchesOSType(t *testing.T) { + fm := NewFingerprintManager() + + for i := 0; i < 20; i++ { + fp := fm.GetFingerprint("token" + string(rune('a'+i))) + validVersions := osVersions[fp.OSType] + found := false + for _, v := range validVersions { + if v == fp.OSVersion { + found = true + break + } + } + if !found { + t.Errorf("OS version %s not valid for OS type %s", fp.OSVersion, fp.OSType) + } + } +} + +func TestFingerprintManager_ConcurrentAccess(t *testing.T) { + fm := NewFingerprintManager() + const numGoroutines = 100 + const numOperations = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < numOperations; j++ { + tokenKey := "token" + string(rune('a'+id%26)) + switch j % 4 { + case 0: + fm.GetFingerprint(tokenKey) + case 1: + fm.Count() + case 2: + fp := fm.GetFingerprint(tokenKey) + req, _ := http.NewRequest("GET", "http://example.com", nil) + fp.ApplyToRequest(req) + case 3: + fm.RemoveFingerprint(tokenKey) + } + } + }(i) + } + + wg.Wait() +} + +func TestKiroHashUniqueness(t *testing.T) { + fm := NewFingerprintManager() + hashes := make(map[string]bool) + + for i := 0; i < 100; i++ { + fp := fm.GetFingerprint("token" + string(rune(i))) + if hashes[fp.KiroHash] { + t.Errorf("duplicate KiroHash detected: %s", fp.KiroHash) + } + hashes[fp.KiroHash] = true + } +} + +func TestKiroHashFormat(t *testing.T) { + fm := NewFingerprintManager() + fp := fm.GetFingerprint("token1") + + if len(fp.KiroHash) != 64 { + t.Errorf("expected KiroHash length 64 (SHA256 hex), got %d", len(fp.KiroHash)) + } + + for _, c := range fp.KiroHash { + if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) { + t.Errorf("invalid hex character in KiroHash: %c", c) + } + } +} diff --git a/internal/auth/kiro/jitter.go b/internal/auth/kiro/jitter.go new file mode 100644 index 00000000..0569a8fb --- /dev/null +++ b/internal/auth/kiro/jitter.go @@ -0,0 +1,174 @@ +package kiro + +import ( + "math/rand" + "sync" + "time" +) + +// Jitter configuration constants +const ( + // JitterPercent is the default percentage of jitter to apply (±30%) + JitterPercent = 0.30 + + // Human-like delay ranges + ShortDelayMin = 50 * time.Millisecond // Minimum for rapid consecutive operations + ShortDelayMax = 200 * time.Millisecond // Maximum for rapid consecutive operations + NormalDelayMin = 1 * time.Second // Minimum for normal thinking time + NormalDelayMax = 3 * time.Second // Maximum for normal thinking time + LongDelayMin = 5 * time.Second // Minimum for reading/resting + LongDelayMax = 10 * time.Second // Maximum for reading/resting + + // Probability thresholds for human-like behavior + ShortDelayProbability = 0.20 // 20% chance of short delay (consecutive ops) + LongDelayProbability = 0.05 // 5% chance of long delay (reading/resting) + NormalDelayProbability = 0.75 // 75% chance of normal delay (thinking) +) + +var ( + jitterRand *rand.Rand + jitterRandOnce sync.Once + jitterMu sync.Mutex + lastRequestTime time.Time +) + +// initJitterRand initializes the random number generator for jitter calculations. +// Uses a time-based seed for unpredictable but reproducible randomness. +func initJitterRand() { + jitterRandOnce.Do(func() { + jitterRand = rand.New(rand.NewSource(time.Now().UnixNano())) + }) +} + +// RandomDelay generates a random delay between min and max duration. +// Thread-safe implementation using mutex protection. +func RandomDelay(min, max time.Duration) time.Duration { + initJitterRand() + jitterMu.Lock() + defer jitterMu.Unlock() + + if min >= max { + return min + } + + rangeMs := max.Milliseconds() - min.Milliseconds() + randomMs := jitterRand.Int63n(rangeMs) + return min + time.Duration(randomMs)*time.Millisecond +} + +// JitterDelay adds jitter to a base delay. +// Applies ±jitterPercent variation to the base delay. +// For example, JitterDelay(1*time.Second, 0.30) returns a value between 700ms and 1300ms. +func JitterDelay(baseDelay time.Duration, jitterPercent float64) time.Duration { + initJitterRand() + jitterMu.Lock() + defer jitterMu.Unlock() + + if jitterPercent <= 0 || jitterPercent > 1 { + jitterPercent = JitterPercent + } + + // Calculate jitter range: base * jitterPercent + jitterRange := float64(baseDelay) * jitterPercent + + // Generate random value in range [-jitterRange, +jitterRange] + jitter := (jitterRand.Float64()*2 - 1) * jitterRange + + result := time.Duration(float64(baseDelay) + jitter) + if result < 0 { + return 0 + } + return result +} + +// JitterDelayDefault applies the default ±30% jitter to a base delay. +func JitterDelayDefault(baseDelay time.Duration) time.Duration { + return JitterDelay(baseDelay, JitterPercent) +} + +// HumanLikeDelay generates a delay that mimics human behavior patterns. +// The delay is selected based on probability distribution: +// - 20% chance: Short delay (50-200ms) - simulates consecutive rapid operations +// - 75% chance: Normal delay (1-3s) - simulates thinking/reading time +// - 5% chance: Long delay (5-10s) - simulates breaks/reading longer content +// +// Returns the delay duration (caller should call time.Sleep with this value). +func HumanLikeDelay() time.Duration { + initJitterRand() + jitterMu.Lock() + defer jitterMu.Unlock() + + // Track time since last request for adaptive behavior + now := time.Now() + timeSinceLastRequest := now.Sub(lastRequestTime) + lastRequestTime = now + + // If requests are very close together, use short delay + if timeSinceLastRequest < 500*time.Millisecond && timeSinceLastRequest > 0 { + rangeMs := ShortDelayMax.Milliseconds() - ShortDelayMin.Milliseconds() + randomMs := jitterRand.Int63n(rangeMs) + return ShortDelayMin + time.Duration(randomMs)*time.Millisecond + } + + // Otherwise, use probability-based selection + roll := jitterRand.Float64() + + var min, max time.Duration + switch { + case roll < ShortDelayProbability: + // Short delay - consecutive operations + min, max = ShortDelayMin, ShortDelayMax + case roll < ShortDelayProbability+LongDelayProbability: + // Long delay - reading/resting + min, max = LongDelayMin, LongDelayMax + default: + // Normal delay - thinking time + min, max = NormalDelayMin, NormalDelayMax + } + + rangeMs := max.Milliseconds() - min.Milliseconds() + randomMs := jitterRand.Int63n(rangeMs) + return min + time.Duration(randomMs)*time.Millisecond +} + +// ApplyHumanLikeDelay applies human-like delay by sleeping. +// This is a convenience function that combines HumanLikeDelay with time.Sleep. +func ApplyHumanLikeDelay() { + delay := HumanLikeDelay() + if delay > 0 { + time.Sleep(delay) + } +} + +// ExponentialBackoffWithJitter calculates retry delay using exponential backoff with jitter. +// Formula: min(baseDelay * 2^attempt + jitter, maxDelay) +// This helps prevent thundering herd problem when multiple clients retry simultaneously. +func ExponentialBackoffWithJitter(attempt int, baseDelay, maxDelay time.Duration) time.Duration { + if attempt < 0 { + attempt = 0 + } + + // Calculate exponential backoff: baseDelay * 2^attempt + backoff := baseDelay * time.Duration(1< maxDelay { + backoff = maxDelay + } + + // Add ±30% jitter + return JitterDelay(backoff, JitterPercent) +} + +// ShouldSkipDelay determines if delay should be skipped based on context. +// Returns true for streaming responses, WebSocket connections, etc. +// This function can be extended to check additional skip conditions. +func ShouldSkipDelay(isStreaming bool) bool { + return isStreaming +} + +// ResetLastRequestTime resets the last request time tracker. +// Useful for testing or when starting a new session. +func ResetLastRequestTime() { + jitterMu.Lock() + defer jitterMu.Unlock() + lastRequestTime = time.Time{} +} diff --git a/internal/auth/kiro/metrics.go b/internal/auth/kiro/metrics.go new file mode 100644 index 00000000..0fe2d0c6 --- /dev/null +++ b/internal/auth/kiro/metrics.go @@ -0,0 +1,187 @@ +package kiro + +import ( + "math" + "sync" + "time" +) + +// TokenMetrics holds performance metrics for a single token. +type TokenMetrics struct { + SuccessRate float64 // Success rate (0.0 - 1.0) + AvgLatency float64 // Average latency in milliseconds + QuotaRemaining float64 // Remaining quota (0.0 - 1.0) + LastUsed time.Time // Last usage timestamp + FailCount int // Consecutive failure count + TotalRequests int // Total request count + successCount int // Internal: successful request count + totalLatency float64 // Internal: cumulative latency +} + +// TokenScorer manages token metrics and scoring. +type TokenScorer struct { + mu sync.RWMutex + metrics map[string]*TokenMetrics + + // Scoring weights + successRateWeight float64 + quotaWeight float64 + latencyWeight float64 + lastUsedWeight float64 + failPenaltyMultiplier float64 +} + +// NewTokenScorer creates a new TokenScorer with default weights. +func NewTokenScorer() *TokenScorer { + return &TokenScorer{ + metrics: make(map[string]*TokenMetrics), + successRateWeight: 0.4, + quotaWeight: 0.25, + latencyWeight: 0.2, + lastUsedWeight: 0.15, + failPenaltyMultiplier: 0.1, + } +} + +// getOrCreateMetrics returns existing metrics or creates new ones. +func (s *TokenScorer) getOrCreateMetrics(tokenKey string) *TokenMetrics { + if m, ok := s.metrics[tokenKey]; ok { + return m + } + m := &TokenMetrics{ + SuccessRate: 1.0, + QuotaRemaining: 1.0, + } + s.metrics[tokenKey] = m + return m +} + +// RecordRequest records the result of a request for a token. +func (s *TokenScorer) RecordRequest(tokenKey string, success bool, latency time.Duration) { + s.mu.Lock() + defer s.mu.Unlock() + + m := s.getOrCreateMetrics(tokenKey) + m.TotalRequests++ + m.LastUsed = time.Now() + m.totalLatency += float64(latency.Milliseconds()) + + if success { + m.successCount++ + m.FailCount = 0 + } else { + m.FailCount++ + } + + // Update derived metrics + if m.TotalRequests > 0 { + m.SuccessRate = float64(m.successCount) / float64(m.TotalRequests) + m.AvgLatency = m.totalLatency / float64(m.TotalRequests) + } +} + +// SetQuotaRemaining updates the remaining quota for a token. +func (s *TokenScorer) SetQuotaRemaining(tokenKey string, quota float64) { + s.mu.Lock() + defer s.mu.Unlock() + + m := s.getOrCreateMetrics(tokenKey) + m.QuotaRemaining = quota +} + +// GetMetrics returns a copy of the metrics for a token. +func (s *TokenScorer) GetMetrics(tokenKey string) *TokenMetrics { + s.mu.RLock() + defer s.mu.RUnlock() + + if m, ok := s.metrics[tokenKey]; ok { + copy := *m + return © + } + return nil +} + +// CalculateScore computes the score for a token (higher is better). +func (s *TokenScorer) CalculateScore(tokenKey string) float64 { + s.mu.RLock() + defer s.mu.RUnlock() + + m, ok := s.metrics[tokenKey] + if !ok { + return 1.0 // New tokens get a high initial score + } + + // Success rate component (0-1) + successScore := m.SuccessRate + + // Quota component (0-1) + quotaScore := m.QuotaRemaining + + // Latency component (normalized, lower is better) + // Using exponential decay: score = e^(-latency/1000) + // 1000ms latency -> ~0.37 score, 100ms -> ~0.90 score + latencyScore := math.Exp(-m.AvgLatency / 1000.0) + if m.TotalRequests == 0 { + latencyScore = 1.0 + } + + // Last used component (prefer tokens not recently used) + // Score increases as time since last use increases + timeSinceUse := time.Since(m.LastUsed).Seconds() + // Normalize: 60 seconds -> ~0.63 score, 0 seconds -> 0 score + lastUsedScore := 1.0 - math.Exp(-timeSinceUse/60.0) + if m.LastUsed.IsZero() { + lastUsedScore = 1.0 + } + + // Calculate weighted score + score := s.successRateWeight*successScore + + s.quotaWeight*quotaScore + + s.latencyWeight*latencyScore + + s.lastUsedWeight*lastUsedScore + + // Apply consecutive failure penalty + if m.FailCount > 0 { + penalty := s.failPenaltyMultiplier * float64(m.FailCount) + score = score * math.Max(0, 1.0-penalty) + } + + return score +} + +// SelectBestToken selects the token with the highest score. +func (s *TokenScorer) SelectBestToken(tokens []string) string { + if len(tokens) == 0 { + return "" + } + if len(tokens) == 1 { + return tokens[0] + } + + bestToken := tokens[0] + bestScore := s.CalculateScore(tokens[0]) + + for _, token := range tokens[1:] { + score := s.CalculateScore(token) + if score > bestScore { + bestScore = score + bestToken = token + } + } + + return bestToken +} + +// ResetMetrics clears all metrics for a token. +func (s *TokenScorer) ResetMetrics(tokenKey string) { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.metrics, tokenKey) +} + +// ResetAllMetrics clears all stored metrics. +func (s *TokenScorer) ResetAllMetrics() { + s.mu.Lock() + defer s.mu.Unlock() + s.metrics = make(map[string]*TokenMetrics) +} diff --git a/internal/auth/kiro/metrics_test.go b/internal/auth/kiro/metrics_test.go new file mode 100644 index 00000000..ffe2a876 --- /dev/null +++ b/internal/auth/kiro/metrics_test.go @@ -0,0 +1,301 @@ +package kiro + +import ( + "sync" + "testing" + "time" +) + +func TestNewTokenScorer(t *testing.T) { + s := NewTokenScorer() + if s == nil { + t.Fatal("expected non-nil TokenScorer") + } + if s.metrics == nil { + t.Error("expected non-nil metrics map") + } + if s.successRateWeight != 0.4 { + t.Errorf("expected successRateWeight 0.4, got %f", s.successRateWeight) + } + if s.quotaWeight != 0.25 { + t.Errorf("expected quotaWeight 0.25, got %f", s.quotaWeight) + } +} + +func TestRecordRequest_Success(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + + m := s.GetMetrics("token1") + if m == nil { + t.Fatal("expected non-nil metrics") + } + if m.TotalRequests != 1 { + t.Errorf("expected TotalRequests 1, got %d", m.TotalRequests) + } + if m.SuccessRate != 1.0 { + t.Errorf("expected SuccessRate 1.0, got %f", m.SuccessRate) + } + if m.FailCount != 0 { + t.Errorf("expected FailCount 0, got %d", m.FailCount) + } + if m.AvgLatency != 100 { + t.Errorf("expected AvgLatency 100, got %f", m.AvgLatency) + } +} + +func TestRecordRequest_Failure(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", false, 200*time.Millisecond) + + m := s.GetMetrics("token1") + if m.SuccessRate != 0.0 { + t.Errorf("expected SuccessRate 0.0, got %f", m.SuccessRate) + } + if m.FailCount != 1 { + t.Errorf("expected FailCount 1, got %d", m.FailCount) + } +} + +func TestRecordRequest_MixedResults(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + s.RecordRequest("token1", true, 100*time.Millisecond) + s.RecordRequest("token1", false, 100*time.Millisecond) + s.RecordRequest("token1", true, 100*time.Millisecond) + + m := s.GetMetrics("token1") + if m.TotalRequests != 4 { + t.Errorf("expected TotalRequests 4, got %d", m.TotalRequests) + } + if m.SuccessRate != 0.75 { + t.Errorf("expected SuccessRate 0.75, got %f", m.SuccessRate) + } + if m.FailCount != 0 { + t.Errorf("expected FailCount 0 (reset on success), got %d", m.FailCount) + } +} + +func TestRecordRequest_ConsecutiveFailures(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + s.RecordRequest("token1", false, 100*time.Millisecond) + s.RecordRequest("token1", false, 100*time.Millisecond) + s.RecordRequest("token1", false, 100*time.Millisecond) + + m := s.GetMetrics("token1") + if m.FailCount != 3 { + t.Errorf("expected FailCount 3, got %d", m.FailCount) + } +} + +func TestSetQuotaRemaining(t *testing.T) { + s := NewTokenScorer() + s.SetQuotaRemaining("token1", 0.5) + + m := s.GetMetrics("token1") + if m.QuotaRemaining != 0.5 { + t.Errorf("expected QuotaRemaining 0.5, got %f", m.QuotaRemaining) + } +} + +func TestGetMetrics_NonExistent(t *testing.T) { + s := NewTokenScorer() + m := s.GetMetrics("nonexistent") + if m != nil { + t.Error("expected nil metrics for non-existent token") + } +} + +func TestGetMetrics_ReturnsCopy(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + + m1 := s.GetMetrics("token1") + m1.TotalRequests = 999 + + m2 := s.GetMetrics("token1") + if m2.TotalRequests == 999 { + t.Error("GetMetrics should return a copy") + } +} + +func TestCalculateScore_NewToken(t *testing.T) { + s := NewTokenScorer() + score := s.CalculateScore("newtoken") + if score != 1.0 { + t.Errorf("expected score 1.0 for new token, got %f", score) + } +} + +func TestCalculateScore_PerfectToken(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 50*time.Millisecond) + s.SetQuotaRemaining("token1", 1.0) + + time.Sleep(100 * time.Millisecond) + score := s.CalculateScore("token1") + if score < 0.5 || score > 1.0 { + t.Errorf("expected high score for perfect token, got %f", score) + } +} + +func TestCalculateScore_FailedToken(t *testing.T) { + s := NewTokenScorer() + for i := 0; i < 5; i++ { + s.RecordRequest("token1", false, 1000*time.Millisecond) + } + s.SetQuotaRemaining("token1", 0.1) + + score := s.CalculateScore("token1") + if score > 0.5 { + t.Errorf("expected low score for failed token, got %f", score) + } +} + +func TestCalculateScore_FailPenalty(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + scoreNoFail := s.CalculateScore("token1") + + s.RecordRequest("token1", false, 100*time.Millisecond) + s.RecordRequest("token1", false, 100*time.Millisecond) + scoreWithFail := s.CalculateScore("token1") + + if scoreWithFail >= scoreNoFail { + t.Errorf("expected lower score with consecutive failures: noFail=%f, withFail=%f", scoreNoFail, scoreWithFail) + } +} + +func TestSelectBestToken_Empty(t *testing.T) { + s := NewTokenScorer() + best := s.SelectBestToken([]string{}) + if best != "" { + t.Errorf("expected empty string for empty tokens, got %s", best) + } +} + +func TestSelectBestToken_SingleToken(t *testing.T) { + s := NewTokenScorer() + best := s.SelectBestToken([]string{"token1"}) + if best != "token1" { + t.Errorf("expected token1, got %s", best) + } +} + +func TestSelectBestToken_MultipleTokens(t *testing.T) { + s := NewTokenScorer() + + s.RecordRequest("bad", false, 1000*time.Millisecond) + s.RecordRequest("bad", false, 1000*time.Millisecond) + s.SetQuotaRemaining("bad", 0.1) + + s.RecordRequest("good", true, 50*time.Millisecond) + s.SetQuotaRemaining("good", 0.9) + + time.Sleep(50 * time.Millisecond) + + best := s.SelectBestToken([]string{"bad", "good"}) + if best != "good" { + t.Errorf("expected good token to be selected, got %s", best) + } +} + +func TestResetMetrics(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + s.ResetMetrics("token1") + + m := s.GetMetrics("token1") + if m != nil { + t.Error("expected nil metrics after reset") + } +} + +func TestResetAllMetrics(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + s.RecordRequest("token2", true, 100*time.Millisecond) + s.RecordRequest("token3", true, 100*time.Millisecond) + + s.ResetAllMetrics() + + if s.GetMetrics("token1") != nil { + t.Error("expected nil metrics for token1 after reset all") + } + if s.GetMetrics("token2") != nil { + t.Error("expected nil metrics for token2 after reset all") + } +} + +func TestTokenScorer_ConcurrentAccess(t *testing.T) { + s := NewTokenScorer() + const numGoroutines = 50 + const numOperations = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + tokenKey := "token" + string(rune('a'+id%10)) + for j := 0; j < numOperations; j++ { + switch j % 6 { + case 0: + s.RecordRequest(tokenKey, j%2 == 0, time.Duration(j)*time.Millisecond) + case 1: + s.SetQuotaRemaining(tokenKey, float64(j%100)/100) + case 2: + s.GetMetrics(tokenKey) + case 3: + s.CalculateScore(tokenKey) + case 4: + s.SelectBestToken([]string{tokenKey, "token_x", "token_y"}) + case 5: + if j%20 == 0 { + s.ResetMetrics(tokenKey) + } + } + } + }(i) + } + + wg.Wait() +} + +func TestAvgLatencyCalculation(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + s.RecordRequest("token1", true, 200*time.Millisecond) + s.RecordRequest("token1", true, 300*time.Millisecond) + + m := s.GetMetrics("token1") + if m.AvgLatency != 200 { + t.Errorf("expected AvgLatency 200, got %f", m.AvgLatency) + } +} + +func TestLastUsedUpdated(t *testing.T) { + s := NewTokenScorer() + before := time.Now() + s.RecordRequest("token1", true, 100*time.Millisecond) + + m := s.GetMetrics("token1") + if m.LastUsed.Before(before) { + t.Error("expected LastUsed to be after test start time") + } + if m.LastUsed.After(time.Now()) { + t.Error("expected LastUsed to be before or equal to now") + } +} + +func TestDefaultQuotaForNewToken(t *testing.T) { + s := NewTokenScorer() + s.RecordRequest("token1", true, 100*time.Millisecond) + + m := s.GetMetrics("token1") + if m.QuotaRemaining != 1.0 { + t.Errorf("expected default QuotaRemaining 1.0, got %f", m.QuotaRemaining) + } +} diff --git a/internal/auth/kiro/oauth.go b/internal/auth/kiro/oauth.go index a7d3eb9a..0609610f 100644 --- a/internal/auth/kiro/oauth.go +++ b/internal/auth/kiro/oauth.go @@ -227,6 +227,7 @@ func (o *KiroOAuth) exchangeCodeForToken(ctx context.Context, code, codeVerifier ExpiresAt: expiresAt.Format(time.RFC3339), AuthMethod: "social", Provider: "", // Caller should preserve original provider + Region: "us-east-1", }, nil } @@ -285,6 +286,7 @@ func (o *KiroOAuth) RefreshToken(ctx context.Context, refreshToken string) (*Kir ExpiresAt: expiresAt.Format(time.RFC3339), AuthMethod: "social", Provider: "", // Caller should preserve original provider + Region: "us-east-1", }, nil } diff --git a/internal/auth/kiro/oauth_web.go b/internal/auth/kiro/oauth_web.go new file mode 100644 index 00000000..13198516 --- /dev/null +++ b/internal/auth/kiro/oauth_web.go @@ -0,0 +1,825 @@ +// Package kiro provides OAuth Web authentication for Kiro. +package kiro + +import ( + "context" + "crypto/rand" + "encoding/base64" + "fmt" + "html/template" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + "github.com/router-for-me/CLIProxyAPI/v6/internal/util" + log "github.com/sirupsen/logrus" +) + +const ( + defaultSessionExpiry = 10 * time.Minute + pollIntervalSeconds = 5 +) + +type authSessionStatus string + +const ( + statusPending authSessionStatus = "pending" + statusSuccess authSessionStatus = "success" + statusFailed authSessionStatus = "failed" +) + +type webAuthSession struct { + stateID string + deviceCode string + userCode string + authURL string + verificationURI string + expiresIn int + interval int + status authSessionStatus + startedAt time.Time + completedAt time.Time + expiresAt time.Time + error string + tokenData *KiroTokenData + ssoClient *SSOOIDCClient + clientID string + clientSecret string + region string + cancelFunc context.CancelFunc + authMethod string // "google", "github", "builder-id", "idc" + startURL string // Used for IDC + codeVerifier string // Used for social auth PKCE + codeChallenge string // Used for social auth PKCE +} + +type OAuthWebHandler struct { + cfg *config.Config + sessions map[string]*webAuthSession + mu sync.RWMutex + onTokenObtained func(*KiroTokenData) +} + +func NewOAuthWebHandler(cfg *config.Config) *OAuthWebHandler { + return &OAuthWebHandler{ + cfg: cfg, + sessions: make(map[string]*webAuthSession), + } +} + +func (h *OAuthWebHandler) SetTokenCallback(callback func(*KiroTokenData)) { + h.onTokenObtained = callback +} + +func (h *OAuthWebHandler) RegisterRoutes(router gin.IRouter) { + oauth := router.Group("/v0/oauth/kiro") + { + oauth.GET("", h.handleSelect) + oauth.GET("/start", h.handleStart) + oauth.GET("/callback", h.handleCallback) + oauth.GET("/social/callback", h.handleSocialCallback) + oauth.GET("/status", h.handleStatus) + oauth.POST("/import", h.handleImportToken) + } +} + +func generateStateID() (string, error) { + b := make([]byte, 16) + if _, err := rand.Read(b); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(b), nil +} + +func (h *OAuthWebHandler) handleSelect(c *gin.Context) { + h.renderSelectPage(c) +} + +func (h *OAuthWebHandler) handleStart(c *gin.Context) { + method := c.Query("method") + + if method == "" { + c.Redirect(http.StatusFound, "/v0/oauth/kiro") + return + } + + switch method { + case "google", "github": + // Google/GitHub social login is not supported for third-party apps + // due to AWS Cognito redirect_uri restrictions + h.renderError(c, "Google/GitHub login is not available for third-party applications. Please use AWS Builder ID or import your token from Kiro IDE.") + case "builder-id": + h.startBuilderIDAuth(c) + case "idc": + h.startIDCAuth(c) + default: + h.renderError(c, fmt.Sprintf("Unknown authentication method: %s", method)) + } +} + +func (h *OAuthWebHandler) startSocialAuth(c *gin.Context, method string) { + stateID, err := generateStateID() + if err != nil { + h.renderError(c, "Failed to generate state parameter") + return + } + + codeVerifier, codeChallenge, err := generatePKCE() + if err != nil { + h.renderError(c, "Failed to generate PKCE parameters") + return + } + + socialClient := NewSocialAuthClient(h.cfg) + + var provider string + if method == "google" { + provider = string(ProviderGoogle) + } else { + provider = string(ProviderGitHub) + } + + redirectURI := h.getSocialCallbackURL(c) + authURL := socialClient.buildLoginURL(provider, redirectURI, codeChallenge, stateID) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + + session := &webAuthSession{ + stateID: stateID, + authMethod: method, + authURL: authURL, + status: statusPending, + startedAt: time.Now(), + expiresIn: 600, + codeVerifier: codeVerifier, + codeChallenge: codeChallenge, + region: "us-east-1", + cancelFunc: cancel, + } + + h.mu.Lock() + h.sessions[stateID] = session + h.mu.Unlock() + + go func() { + <-ctx.Done() + h.mu.Lock() + if session.status == statusPending { + session.status = statusFailed + session.error = "Authentication timed out" + } + h.mu.Unlock() + }() + + c.Redirect(http.StatusFound, authURL) +} + +func (h *OAuthWebHandler) getSocialCallbackURL(c *gin.Context) string { + scheme := "http" + if c.Request.TLS != nil || c.GetHeader("X-Forwarded-Proto") == "https" { + scheme = "https" + } + return fmt.Sprintf("%s://%s/v0/oauth/kiro/social/callback", scheme, c.Request.Host) +} + +func (h *OAuthWebHandler) startBuilderIDAuth(c *gin.Context) { + stateID, err := generateStateID() + if err != nil { + h.renderError(c, "Failed to generate state parameter") + return + } + + region := defaultIDCRegion + startURL := builderIDStartURL + + ssoClient := NewSSOOIDCClient(h.cfg) + + regResp, err := ssoClient.RegisterClientWithRegion(c.Request.Context(), region) + if err != nil { + log.Errorf("OAuth Web: failed to register client: %v", err) + h.renderError(c, fmt.Sprintf("Failed to register client: %v", err)) + return + } + + authResp, err := ssoClient.StartDeviceAuthorizationWithIDC( + c.Request.Context(), + regResp.ClientID, + regResp.ClientSecret, + startURL, + region, + ) + if err != nil { + log.Errorf("OAuth Web: failed to start device authorization: %v", err) + h.renderError(c, fmt.Sprintf("Failed to start device authorization: %v", err)) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(authResp.ExpiresIn)*time.Second) + + session := &webAuthSession{ + stateID: stateID, + deviceCode: authResp.DeviceCode, + userCode: authResp.UserCode, + authURL: authResp.VerificationURIComplete, + verificationURI: authResp.VerificationURI, + expiresIn: authResp.ExpiresIn, + interval: authResp.Interval, + status: statusPending, + startedAt: time.Now(), + ssoClient: ssoClient, + clientID: regResp.ClientID, + clientSecret: regResp.ClientSecret, + region: region, + authMethod: "builder-id", + startURL: startURL, + cancelFunc: cancel, + } + + h.mu.Lock() + h.sessions[stateID] = session + h.mu.Unlock() + + go h.pollForToken(ctx, session) + + h.renderStartPage(c, session) +} + +func (h *OAuthWebHandler) startIDCAuth(c *gin.Context) { + startURL := c.Query("startUrl") + region := c.Query("region") + + if startURL == "" { + h.renderError(c, "Missing startUrl parameter for IDC authentication") + return + } + if region == "" { + region = defaultIDCRegion + } + + stateID, err := generateStateID() + if err != nil { + h.renderError(c, "Failed to generate state parameter") + return + } + + ssoClient := NewSSOOIDCClient(h.cfg) + + regResp, err := ssoClient.RegisterClientWithRegion(c.Request.Context(), region) + if err != nil { + log.Errorf("OAuth Web: failed to register client: %v", err) + h.renderError(c, fmt.Sprintf("Failed to register client: %v", err)) + return + } + + authResp, err := ssoClient.StartDeviceAuthorizationWithIDC( + c.Request.Context(), + regResp.ClientID, + regResp.ClientSecret, + startURL, + region, + ) + if err != nil { + log.Errorf("OAuth Web: failed to start device authorization: %v", err) + h.renderError(c, fmt.Sprintf("Failed to start device authorization: %v", err)) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(authResp.ExpiresIn)*time.Second) + + session := &webAuthSession{ + stateID: stateID, + deviceCode: authResp.DeviceCode, + userCode: authResp.UserCode, + authURL: authResp.VerificationURIComplete, + verificationURI: authResp.VerificationURI, + expiresIn: authResp.ExpiresIn, + interval: authResp.Interval, + status: statusPending, + startedAt: time.Now(), + ssoClient: ssoClient, + clientID: regResp.ClientID, + clientSecret: regResp.ClientSecret, + region: region, + authMethod: "idc", + startURL: startURL, + cancelFunc: cancel, + } + + h.mu.Lock() + h.sessions[stateID] = session + h.mu.Unlock() + + go h.pollForToken(ctx, session) + + h.renderStartPage(c, session) +} + +func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSession) { + defer session.cancelFunc() + + interval := time.Duration(session.interval) * time.Second + if interval < time.Duration(pollIntervalSeconds)*time.Second { + interval = time.Duration(pollIntervalSeconds) * time.Second + } + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + h.mu.Lock() + if session.status == statusPending { + session.status = statusFailed + session.error = "Authentication timed out" + } + h.mu.Unlock() + return + case <-ticker.C: + tokenResp, err := h.ssoClient(session).CreateTokenWithRegion( + ctx, + session.clientID, + session.clientSecret, + session.deviceCode, + session.region, + ) + + if err != nil { + errStr := err.Error() + if errStr == ErrAuthorizationPending.Error() { + continue + } + if errStr == ErrSlowDown.Error() { + interval += 5 * time.Second + ticker.Reset(interval) + continue + } + + h.mu.Lock() + session.status = statusFailed + session.error = errStr + session.completedAt = time.Now() + h.mu.Unlock() + + log.Errorf("OAuth Web: token polling failed: %v", err) + return + } + + expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) + profileArn := session.ssoClient.fetchProfileArn(ctx, tokenResp.AccessToken) + email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken) + + tokenData := &KiroTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: profileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: session.authMethod, + Provider: "AWS", + ClientID: session.clientID, + ClientSecret: session.clientSecret, + Email: email, + Region: session.region, + } + + h.mu.Lock() + session.status = statusSuccess + session.completedAt = time.Now() + session.expiresAt = expiresAt + session.tokenData = tokenData + h.mu.Unlock() + + if h.onTokenObtained != nil { + h.onTokenObtained(tokenData) + } + + // Save token to file + h.saveTokenToFile(tokenData) + + log.Infof("OAuth Web: authentication successful for %s", email) + return + } + } +} + +// saveTokenToFile saves the token data to the auth directory +func (h *OAuthWebHandler) saveTokenToFile(tokenData *KiroTokenData) { + // Get auth directory from config or use default + authDir := "" + if h.cfg != nil && h.cfg.AuthDir != "" { + var err error + authDir, err = util.ResolveAuthDir(h.cfg.AuthDir) + if err != nil { + log.Errorf("OAuth Web: failed to resolve auth directory: %v", err) + } + } + + // Fall back to default location + if authDir == "" { + home, err := os.UserHomeDir() + if err != nil { + log.Errorf("OAuth Web: failed to get home directory: %v", err) + return + } + authDir = filepath.Join(home, ".cli-proxy-api") + } + + // Create directory if not exists + if err := os.MkdirAll(authDir, 0700); err != nil { + log.Errorf("OAuth Web: failed to create auth directory: %v", err) + return + } + + // Generate filename based on auth method + // Format: kiro-{authMethod}.json or kiro-{authMethod}-{email}.json + fileName := fmt.Sprintf("kiro-%s.json", tokenData.AuthMethod) + if tokenData.Email != "" { + // Sanitize email for filename (replace @ and . with -) + sanitizedEmail := tokenData.Email + sanitizedEmail = strings.ReplaceAll(sanitizedEmail, "@", "-") + sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-") + fileName = fmt.Sprintf("kiro-%s-%s.json", tokenData.AuthMethod, sanitizedEmail) + } + + authFilePath := filepath.Join(authDir, fileName) + + // Convert to storage format and save + storage := &KiroTokenStorage{ + Type: "kiro", + AccessToken: tokenData.AccessToken, + RefreshToken: tokenData.RefreshToken, + ProfileArn: tokenData.ProfileArn, + ExpiresAt: tokenData.ExpiresAt, + AuthMethod: tokenData.AuthMethod, + Provider: tokenData.Provider, + LastRefresh: time.Now().Format(time.RFC3339), + ClientID: tokenData.ClientID, + ClientSecret: tokenData.ClientSecret, + Region: tokenData.Region, + StartURL: tokenData.StartURL, + Email: tokenData.Email, + } + + if err := storage.SaveTokenToFile(authFilePath); err != nil { + log.Errorf("OAuth Web: failed to save token to file: %v", err) + return + } + + log.Infof("OAuth Web: token saved to %s", authFilePath) +} + +func (h *OAuthWebHandler) ssoClient(session *webAuthSession) *SSOOIDCClient { + return session.ssoClient +} + +func (h *OAuthWebHandler) handleCallback(c *gin.Context) { + stateID := c.Query("state") + errParam := c.Query("error") + + if errParam != "" { + h.renderError(c, errParam) + return + } + + if stateID == "" { + h.renderError(c, "Missing state parameter") + return + } + + h.mu.RLock() + session, exists := h.sessions[stateID] + h.mu.RUnlock() + + if !exists { + h.renderError(c, "Invalid or expired session") + return + } + + if session.status == statusSuccess { + h.renderSuccess(c, session) + } else if session.status == statusFailed { + h.renderError(c, session.error) + } else { + c.Redirect(http.StatusFound, "/v0/oauth/kiro/start") + } +} + +func (h *OAuthWebHandler) handleSocialCallback(c *gin.Context) { + stateID := c.Query("state") + code := c.Query("code") + errParam := c.Query("error") + + if errParam != "" { + h.renderError(c, errParam) + return + } + + if stateID == "" { + h.renderError(c, "Missing state parameter") + return + } + + if code == "" { + h.renderError(c, "Missing authorization code") + return + } + + h.mu.RLock() + session, exists := h.sessions[stateID] + h.mu.RUnlock() + + if !exists { + h.renderError(c, "Invalid or expired session") + return + } + + if session.authMethod != "google" && session.authMethod != "github" { + h.renderError(c, "Invalid session type for social callback") + return + } + + socialClient := NewSocialAuthClient(h.cfg) + redirectURI := h.getSocialCallbackURL(c) + + tokenReq := &CreateTokenRequest{ + Code: code, + CodeVerifier: session.codeVerifier, + RedirectURI: redirectURI, + } + + tokenResp, err := socialClient.CreateToken(c.Request.Context(), tokenReq) + if err != nil { + log.Errorf("OAuth Web: social token exchange failed: %v", err) + h.mu.Lock() + session.status = statusFailed + session.error = fmt.Sprintf("Token exchange failed: %v", err) + session.completedAt = time.Now() + h.mu.Unlock() + h.renderError(c, session.error) + return + } + + expiresIn := tokenResp.ExpiresIn + if expiresIn <= 0 { + expiresIn = 3600 + } + expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second) + + email := ExtractEmailFromJWT(tokenResp.AccessToken) + + var provider string + if session.authMethod == "google" { + provider = string(ProviderGoogle) + } else { + provider = string(ProviderGitHub) + } + + tokenData := &KiroTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: tokenResp.ProfileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: session.authMethod, + Provider: provider, + Email: email, + Region: "us-east-1", + } + + h.mu.Lock() + session.status = statusSuccess + session.completedAt = time.Now() + session.expiresAt = expiresAt + session.tokenData = tokenData + h.mu.Unlock() + + if session.cancelFunc != nil { + session.cancelFunc() + } + + if h.onTokenObtained != nil { + h.onTokenObtained(tokenData) + } + + // Save token to file + h.saveTokenToFile(tokenData) + + log.Infof("OAuth Web: social authentication successful for %s via %s", email, provider) + h.renderSuccess(c, session) +} + +func (h *OAuthWebHandler) handleStatus(c *gin.Context) { + stateID := c.Query("state") + if stateID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "missing state parameter"}) + return + } + + h.mu.RLock() + session, exists := h.sessions[stateID] + h.mu.RUnlock() + + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) + return + } + + response := gin.H{ + "status": string(session.status), + } + + switch session.status { + case statusPending: + elapsed := time.Since(session.startedAt).Seconds() + remaining := float64(session.expiresIn) - elapsed + if remaining < 0 { + remaining = 0 + } + response["remaining_seconds"] = int(remaining) + case statusSuccess: + response["completed_at"] = session.completedAt.Format(time.RFC3339) + response["expires_at"] = session.expiresAt.Format(time.RFC3339) + case statusFailed: + response["error"] = session.error + response["failed_at"] = session.completedAt.Format(time.RFC3339) + } + + c.JSON(http.StatusOK, response) +} + +func (h *OAuthWebHandler) renderStartPage(c *gin.Context, session *webAuthSession) { + tmpl, err := template.New("start").Parse(oauthWebStartPageHTML) + if err != nil { + log.Errorf("OAuth Web: failed to parse template: %v", err) + c.String(http.StatusInternalServerError, "Template error") + return + } + + data := map[string]interface{}{ + "AuthURL": session.authURL, + "UserCode": session.userCode, + "ExpiresIn": session.expiresIn, + "StateID": session.stateID, + } + + c.Header("Content-Type", "text/html; charset=utf-8") + if err := tmpl.Execute(c.Writer, data); err != nil { + log.Errorf("OAuth Web: failed to render template: %v", err) + } +} + +func (h *OAuthWebHandler) renderSelectPage(c *gin.Context) { + tmpl, err := template.New("select").Parse(oauthWebSelectPageHTML) + if err != nil { + log.Errorf("OAuth Web: failed to parse select template: %v", err) + c.String(http.StatusInternalServerError, "Template error") + return + } + + c.Header("Content-Type", "text/html; charset=utf-8") + if err := tmpl.Execute(c.Writer, nil); err != nil { + log.Errorf("OAuth Web: failed to render select template: %v", err) + } +} + +func (h *OAuthWebHandler) renderError(c *gin.Context, errMsg string) { + tmpl, err := template.New("error").Parse(oauthWebErrorPageHTML) + if err != nil { + log.Errorf("OAuth Web: failed to parse error template: %v", err) + c.String(http.StatusInternalServerError, "Template error") + return + } + + data := map[string]interface{}{ + "Error": errMsg, + } + + c.Header("Content-Type", "text/html; charset=utf-8") + c.Status(http.StatusBadRequest) + if err := tmpl.Execute(c.Writer, data); err != nil { + log.Errorf("OAuth Web: failed to render error template: %v", err) + } +} + +func (h *OAuthWebHandler) renderSuccess(c *gin.Context, session *webAuthSession) { + tmpl, err := template.New("success").Parse(oauthWebSuccessPageHTML) + if err != nil { + log.Errorf("OAuth Web: failed to parse success template: %v", err) + c.String(http.StatusInternalServerError, "Template error") + return + } + + data := map[string]interface{}{ + "ExpiresAt": session.expiresAt.Format(time.RFC3339), + } + + c.Header("Content-Type", "text/html; charset=utf-8") + if err := tmpl.Execute(c.Writer, data); err != nil { + log.Errorf("OAuth Web: failed to render success template: %v", err) + } +} + +func (h *OAuthWebHandler) CleanupExpiredSessions() { + h.mu.Lock() + defer h.mu.Unlock() + + now := time.Now() + for id, session := range h.sessions { + if session.status != statusPending && now.Sub(session.completedAt) > 30*time.Minute { + delete(h.sessions, id) + } else if session.status == statusPending && now.Sub(session.startedAt) > defaultSessionExpiry { + session.cancelFunc() + delete(h.sessions, id) + } + } +} + +func (h *OAuthWebHandler) GetSession(stateID string) (*webAuthSession, bool) { + h.mu.RLock() + defer h.mu.RUnlock() + session, exists := h.sessions[stateID] + return session, exists +} + +// ImportTokenRequest represents the request body for token import +type ImportTokenRequest struct { + RefreshToken string `json:"refreshToken"` +} + +// handleImportToken handles manual refresh token import from Kiro IDE +func (h *OAuthWebHandler) handleImportToken(c *gin.Context) { + var req ImportTokenRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "error": "Invalid request body", + }) + return + } + + refreshToken := strings.TrimSpace(req.RefreshToken) + if refreshToken == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "error": "Refresh token is required", + }) + return + } + + // Validate token format + if !strings.HasPrefix(refreshToken, "aorAAAAAG") { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "error": "Invalid token format. Token should start with aorAAAAAG...", + }) + return + } + + // Create social auth client to refresh and validate the token + socialClient := NewSocialAuthClient(h.cfg) + + // Refresh the token to validate it and get access token + tokenData, err := socialClient.RefreshSocialToken(c.Request.Context(), refreshToken) + if err != nil { + log.Errorf("OAuth Web: token refresh failed during import: %v", err) + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "error": fmt.Sprintf("Token validation failed: %v", err), + }) + return + } + + // Set the original refresh token (the refreshed one might be empty) + if tokenData.RefreshToken == "" { + tokenData.RefreshToken = refreshToken + } + tokenData.AuthMethod = "social" + tokenData.Provider = "imported" + + // Notify callback if set + if h.onTokenObtained != nil { + h.onTokenObtained(tokenData) + } + + // Save token to file + h.saveTokenToFile(tokenData) + + // Generate filename for response + fileName := fmt.Sprintf("kiro-%s.json", tokenData.AuthMethod) + if tokenData.Email != "" { + sanitizedEmail := strings.ReplaceAll(tokenData.Email, "@", "-") + sanitizedEmail = strings.ReplaceAll(sanitizedEmail, ".", "-") + fileName = fmt.Sprintf("kiro-%s-%s.json", tokenData.AuthMethod, sanitizedEmail) + } + + log.Infof("OAuth Web: token imported successfully") + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "Token imported successfully", + "fileName": fileName, + }) +} diff --git a/internal/auth/kiro/oauth_web.go.bak b/internal/auth/kiro/oauth_web.go.bak new file mode 100644 index 00000000..22d7809b --- /dev/null +++ b/internal/auth/kiro/oauth_web.go.bak @@ -0,0 +1,385 @@ +// Package kiro provides OAuth Web authentication for Kiro. +package kiro + +import ( + "context" + "crypto/rand" + "encoding/base64" + "fmt" + "html/template" + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + log "github.com/sirupsen/logrus" +) + +const ( + defaultSessionExpiry = 10 * time.Minute + pollIntervalSeconds = 5 +) + +type authSessionStatus string + +const ( + statusPending authSessionStatus = "pending" + statusSuccess authSessionStatus = "success" + statusFailed authSessionStatus = "failed" +) + +type webAuthSession struct { + stateID string + deviceCode string + userCode string + authURL string + verificationURI string + expiresIn int + interval int + status authSessionStatus + startedAt time.Time + completedAt time.Time + expiresAt time.Time + error string + tokenData *KiroTokenData + ssoClient *SSOOIDCClient + clientID string + clientSecret string + region string + cancelFunc context.CancelFunc +} + +type OAuthWebHandler struct { + cfg *config.Config + sessions map[string]*webAuthSession + mu sync.RWMutex + onTokenObtained func(*KiroTokenData) +} + +func NewOAuthWebHandler(cfg *config.Config) *OAuthWebHandler { + return &OAuthWebHandler{ + cfg: cfg, + sessions: make(map[string]*webAuthSession), + } +} + +func (h *OAuthWebHandler) SetTokenCallback(callback func(*KiroTokenData)) { + h.onTokenObtained = callback +} + +func (h *OAuthWebHandler) RegisterRoutes(router gin.IRouter) { + oauth := router.Group("/v0/oauth/kiro") + { + oauth.GET("/start", h.handleStart) + oauth.GET("/callback", h.handleCallback) + oauth.GET("/status", h.handleStatus) + } +} + +func generateStateID() (string, error) { + b := make([]byte, 16) + if _, err := rand.Read(b); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(b), nil +} + +func (h *OAuthWebHandler) handleStart(c *gin.Context) { + stateID, err := generateStateID() + if err != nil { + h.renderError(c, "Failed to generate state parameter") + return + } + + region := defaultIDCRegion + startURL := builderIDStartURL + + ssoClient := NewSSOOIDCClient(h.cfg) + + regResp, err := ssoClient.RegisterClientWithRegion(c.Request.Context(), region) + if err != nil { + log.Errorf("OAuth Web: failed to register client: %v", err) + h.renderError(c, fmt.Sprintf("Failed to register client: %v", err)) + return + } + + authResp, err := ssoClient.StartDeviceAuthorizationWithIDC( + c.Request.Context(), + regResp.ClientID, + regResp.ClientSecret, + startURL, + region, + ) + if err != nil { + log.Errorf("OAuth Web: failed to start device authorization: %v", err) + h.renderError(c, fmt.Sprintf("Failed to start device authorization: %v", err)) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(authResp.ExpiresIn)*time.Second) + + session := &webAuthSession{ + stateID: stateID, + deviceCode: authResp.DeviceCode, + userCode: authResp.UserCode, + authURL: authResp.VerificationURIComplete, + verificationURI: authResp.VerificationURI, + expiresIn: authResp.ExpiresIn, + interval: authResp.Interval, + status: statusPending, + startedAt: time.Now(), + ssoClient: ssoClient, + clientID: regResp.ClientID, + clientSecret: regResp.ClientSecret, + region: region, + cancelFunc: cancel, + } + + h.mu.Lock() + h.sessions[stateID] = session + h.mu.Unlock() + + go h.pollForToken(ctx, session) + + h.renderStartPage(c, session) +} + +func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSession) { + defer session.cancelFunc() + + interval := time.Duration(session.interval) * time.Second + if interval < time.Duration(pollIntervalSeconds)*time.Second { + interval = time.Duration(pollIntervalSeconds) * time.Second + } + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + h.mu.Lock() + if session.status == statusPending { + session.status = statusFailed + session.error = "Authentication timed out" + } + h.mu.Unlock() + return + case <-ticker.C: + tokenResp, err := h.ssoClient(session).CreateTokenWithRegion( + ctx, + session.clientID, + session.clientSecret, + session.deviceCode, + session.region, + ) + + if err != nil { + errStr := err.Error() + if errStr == ErrAuthorizationPending.Error() { + continue + } + if errStr == ErrSlowDown.Error() { + interval += 5 * time.Second + ticker.Reset(interval) + continue + } + + h.mu.Lock() + session.status = statusFailed + session.error = errStr + session.completedAt = time.Now() + h.mu.Unlock() + + log.Errorf("OAuth Web: token polling failed: %v", err) + return + } + + expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) + profileArn := session.ssoClient.fetchProfileArn(ctx, tokenResp.AccessToken) + email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken) + + tokenData := &KiroTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: profileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: "builder-id", + Provider: "AWS", + ClientID: session.clientID, + ClientSecret: session.clientSecret, + Email: email, + } + + h.mu.Lock() + session.status = statusSuccess + session.completedAt = time.Now() + session.expiresAt = expiresAt + session.tokenData = tokenData + h.mu.Unlock() + + if h.onTokenObtained != nil { + h.onTokenObtained(tokenData) + } + + log.Infof("OAuth Web: authentication successful for %s", email) + return + } + } +} + +func (h *OAuthWebHandler) ssoClient(session *webAuthSession) *SSOOIDCClient { + return session.ssoClient +} + +func (h *OAuthWebHandler) handleCallback(c *gin.Context) { + stateID := c.Query("state") + errParam := c.Query("error") + + if errParam != "" { + h.renderError(c, errParam) + return + } + + if stateID == "" { + h.renderError(c, "Missing state parameter") + return + } + + h.mu.RLock() + session, exists := h.sessions[stateID] + h.mu.RUnlock() + + if !exists { + h.renderError(c, "Invalid or expired session") + return + } + + if session.status == statusSuccess { + h.renderSuccess(c, session) + } else if session.status == statusFailed { + h.renderError(c, session.error) + } else { + c.Redirect(http.StatusFound, "/v0/oauth/kiro/start") + } +} + +func (h *OAuthWebHandler) handleStatus(c *gin.Context) { + stateID := c.Query("state") + if stateID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "missing state parameter"}) + return + } + + h.mu.RLock() + session, exists := h.sessions[stateID] + h.mu.RUnlock() + + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) + return + } + + response := gin.H{ + "status": string(session.status), + } + + switch session.status { + case statusPending: + elapsed := time.Since(session.startedAt).Seconds() + remaining := float64(session.expiresIn) - elapsed + if remaining < 0 { + remaining = 0 + } + response["remaining_seconds"] = int(remaining) + case statusSuccess: + response["completed_at"] = session.completedAt.Format(time.RFC3339) + response["expires_at"] = session.expiresAt.Format(time.RFC3339) + case statusFailed: + response["error"] = session.error + response["failed_at"] = session.completedAt.Format(time.RFC3339) + } + + c.JSON(http.StatusOK, response) +} + +func (h *OAuthWebHandler) renderStartPage(c *gin.Context, session *webAuthSession) { + tmpl, err := template.New("start").Parse(oauthWebStartPageHTML) + if err != nil { + log.Errorf("OAuth Web: failed to parse template: %v", err) + c.String(http.StatusInternalServerError, "Template error") + return + } + + data := map[string]interface{}{ + "AuthURL": session.authURL, + "UserCode": session.userCode, + "ExpiresIn": session.expiresIn, + "StateID": session.stateID, + } + + c.Header("Content-Type", "text/html; charset=utf-8") + if err := tmpl.Execute(c.Writer, data); err != nil { + log.Errorf("OAuth Web: failed to render template: %v", err) + } +} + +func (h *OAuthWebHandler) renderError(c *gin.Context, errMsg string) { + tmpl, err := template.New("error").Parse(oauthWebErrorPageHTML) + if err != nil { + log.Errorf("OAuth Web: failed to parse error template: %v", err) + c.String(http.StatusInternalServerError, "Template error") + return + } + + data := map[string]interface{}{ + "Error": errMsg, + } + + c.Header("Content-Type", "text/html; charset=utf-8") + c.Status(http.StatusBadRequest) + if err := tmpl.Execute(c.Writer, data); err != nil { + log.Errorf("OAuth Web: failed to render error template: %v", err) + } +} + +func (h *OAuthWebHandler) renderSuccess(c *gin.Context, session *webAuthSession) { + tmpl, err := template.New("success").Parse(oauthWebSuccessPageHTML) + if err != nil { + log.Errorf("OAuth Web: failed to parse success template: %v", err) + c.String(http.StatusInternalServerError, "Template error") + return + } + + data := map[string]interface{}{ + "ExpiresAt": session.expiresAt.Format(time.RFC3339), + } + + c.Header("Content-Type", "text/html; charset=utf-8") + if err := tmpl.Execute(c.Writer, data); err != nil { + log.Errorf("OAuth Web: failed to render success template: %v", err) + } +} + +func (h *OAuthWebHandler) CleanupExpiredSessions() { + h.mu.Lock() + defer h.mu.Unlock() + + now := time.Now() + for id, session := range h.sessions { + if session.status != statusPending && now.Sub(session.completedAt) > 30*time.Minute { + delete(h.sessions, id) + } else if session.status == statusPending && now.Sub(session.startedAt) > defaultSessionExpiry { + session.cancelFunc() + delete(h.sessions, id) + } + } +} + +func (h *OAuthWebHandler) GetSession(stateID string) (*webAuthSession, bool) { + h.mu.RLock() + defer h.mu.RUnlock() + session, exists := h.sessions[stateID] + return session, exists +} diff --git a/internal/auth/kiro/oauth_web_templates.go b/internal/auth/kiro/oauth_web_templates.go new file mode 100644 index 00000000..064a1ff9 --- /dev/null +++ b/internal/auth/kiro/oauth_web_templates.go @@ -0,0 +1,732 @@ +// Package kiro provides OAuth Web authentication templates. +package kiro + +const ( + oauthWebStartPageHTML = ` + + + + + AWS SSO Authentication + + + +
+

🔐 AWS SSO Authentication

+

Follow the steps below to complete authentication

+ +
+
+ 1 + Click the button below to open the authorization page +
+ + 🚀 Open Authorization Page + +
+ +
+
+ 2 + Enter the verification code below +
+
+
Verification Code
+
{{.UserCode}}
+
+
+ +
+
+ 3 + Complete AWS SSO login +
+

+ Use your AWS SSO account to login and authorize +

+
+ +
+ +
{{.ExpiresIn}}s
+
+ Waiting for authorization... +
+
+ +
+ 💡 Tip: The authorization page will open in a new tab. This page will automatically update once authorization is complete. +
+
+ + + +` + + oauthWebErrorPageHTML = ` + + + + + Authentication Failed + + + +
+

❌ Authentication Failed

+
+

Error:

+

{{.Error}}

+
+ 🔄 Retry +
+ +` + + oauthWebSuccessPageHTML = ` + + + + + Authentication Successful + + + +
+
✅
+

Authentication Successful!

+
+

You can close this window.

+
+
Token expires: {{.ExpiresAt}}
+
+ +` + + oauthWebSelectPageHTML = ` + + + + + Select Authentication Method + + + +
+

🔐 Select Authentication Method

+

Choose how you want to authenticate with Kiro

+ +
+ + 🔶 + AWS Builder ID (Recommended) + + + + +
or
+ + +
+ +
+
+ + +
+ + +
Your AWS Identity Center Start URL
+
+ +
+ + +
AWS Region for your Identity Center
+
+ + +
+
+ +
+
+
+ + +
Copy from Kiro IDE: ~/.kiro/kiro-auth-token.json → refreshToken field
+
+ + + +
+
+
+ +
+ ⚠️ Note: Google and GitHub login are not available for third-party applications due to AWS Cognito restrictions. Please use AWS Builder ID or import your token from Kiro IDE. +
+ +
+ 💡 How to get RefreshToken:
+ 1. Open Kiro IDE and login with Google/GitHub
+ 2. Find the token file: ~/.kiro/kiro-auth-token.json
+ 3. Copy the refreshToken value and paste it above +
+
+ + + +` +) diff --git a/internal/auth/kiro/rate_limiter.go b/internal/auth/kiro/rate_limiter.go new file mode 100644 index 00000000..3c240ebe --- /dev/null +++ b/internal/auth/kiro/rate_limiter.go @@ -0,0 +1,316 @@ +package kiro + +import ( + "math" + "math/rand" + "strings" + "sync" + "time" +) + +const ( + DefaultMinTokenInterval = 10 * time.Second + DefaultMaxTokenInterval = 30 * time.Second + DefaultDailyMaxRequests = 500 + DefaultJitterPercent = 0.3 + DefaultBackoffBase = 2 * time.Minute + DefaultBackoffMax = 60 * time.Minute + DefaultBackoffMultiplier = 2.0 + DefaultSuspendCooldown = 24 * time.Hour +) + +// TokenState Token 状态 +type TokenState struct { + LastRequest time.Time + RequestCount int + CooldownEnd time.Time + FailCount int + DailyRequests int + DailyResetTime time.Time + IsSuspended bool + SuspendedAt time.Time + SuspendReason string +} + +// RateLimiter 频率限制器 +type RateLimiter struct { + mu sync.RWMutex + states map[string]*TokenState + minTokenInterval time.Duration + maxTokenInterval time.Duration + dailyMaxRequests int + jitterPercent float64 + backoffBase time.Duration + backoffMax time.Duration + backoffMultiplier float64 + suspendCooldown time.Duration + rng *rand.Rand +} + +// NewRateLimiter 创建默认配置的频率限制器 +func NewRateLimiter() *RateLimiter { + return &RateLimiter{ + states: make(map[string]*TokenState), + minTokenInterval: DefaultMinTokenInterval, + maxTokenInterval: DefaultMaxTokenInterval, + dailyMaxRequests: DefaultDailyMaxRequests, + jitterPercent: DefaultJitterPercent, + backoffBase: DefaultBackoffBase, + backoffMax: DefaultBackoffMax, + backoffMultiplier: DefaultBackoffMultiplier, + suspendCooldown: DefaultSuspendCooldown, + rng: rand.New(rand.NewSource(time.Now().UnixNano())), + } +} + +// RateLimiterConfig 频率限制器配置 +type RateLimiterConfig struct { + MinTokenInterval time.Duration + MaxTokenInterval time.Duration + DailyMaxRequests int + JitterPercent float64 + BackoffBase time.Duration + BackoffMax time.Duration + BackoffMultiplier float64 + SuspendCooldown time.Duration +} + +// NewRateLimiterWithConfig 使用自定义配置创建频率限制器 +func NewRateLimiterWithConfig(cfg RateLimiterConfig) *RateLimiter { + rl := NewRateLimiter() + if cfg.MinTokenInterval > 0 { + rl.minTokenInterval = cfg.MinTokenInterval + } + if cfg.MaxTokenInterval > 0 { + rl.maxTokenInterval = cfg.MaxTokenInterval + } + if cfg.DailyMaxRequests > 0 { + rl.dailyMaxRequests = cfg.DailyMaxRequests + } + if cfg.JitterPercent > 0 { + rl.jitterPercent = cfg.JitterPercent + } + if cfg.BackoffBase > 0 { + rl.backoffBase = cfg.BackoffBase + } + if cfg.BackoffMax > 0 { + rl.backoffMax = cfg.BackoffMax + } + if cfg.BackoffMultiplier > 0 { + rl.backoffMultiplier = cfg.BackoffMultiplier + } + if cfg.SuspendCooldown > 0 { + rl.suspendCooldown = cfg.SuspendCooldown + } + return rl +} + +// getOrCreateState 获取或创建 Token 状态 +func (rl *RateLimiter) getOrCreateState(tokenKey string) *TokenState { + state, exists := rl.states[tokenKey] + if !exists { + state = &TokenState{ + DailyResetTime: time.Now().Truncate(24 * time.Hour).Add(24 * time.Hour), + } + rl.states[tokenKey] = state + } + return state +} + +// resetDailyIfNeeded 如果需要则重置每日计数 +func (rl *RateLimiter) resetDailyIfNeeded(state *TokenState) { + now := time.Now() + if now.After(state.DailyResetTime) { + state.DailyRequests = 0 + state.DailyResetTime = now.Truncate(24 * time.Hour).Add(24 * time.Hour) + } +} + +// calculateInterval 计算带抖动的随机间隔 +func (rl *RateLimiter) calculateInterval() time.Duration { + baseInterval := rl.minTokenInterval + time.Duration(rl.rng.Int63n(int64(rl.maxTokenInterval-rl.minTokenInterval))) + jitter := time.Duration(float64(baseInterval) * rl.jitterPercent * (rl.rng.Float64()*2 - 1)) + return baseInterval + jitter +} + +// WaitForToken 等待 Token 可用（带抖动的随机间隔） +func (rl *RateLimiter) WaitForToken(tokenKey string) { + rl.mu.Lock() + state := rl.getOrCreateState(tokenKey) + rl.resetDailyIfNeeded(state) + + now := time.Now() + + // 检查是否在冷却期 + if now.Before(state.CooldownEnd) { + waitTime := state.CooldownEnd.Sub(now) + rl.mu.Unlock() + time.Sleep(waitTime) + rl.mu.Lock() + state = rl.getOrCreateState(tokenKey) + now = time.Now() + } + + // 计算距离上次请求的间隔 + interval := rl.calculateInterval() + nextAllowedTime := state.LastRequest.Add(interval) + + if now.Before(nextAllowedTime) { + waitTime := nextAllowedTime.Sub(now) + rl.mu.Unlock() + time.Sleep(waitTime) + rl.mu.Lock() + state = rl.getOrCreateState(tokenKey) + } + + state.LastRequest = time.Now() + state.RequestCount++ + state.DailyRequests++ + rl.mu.Unlock() +} + +// MarkTokenFailed 标记 Token 失败 +func (rl *RateLimiter) MarkTokenFailed(tokenKey string) { + rl.mu.Lock() + defer rl.mu.Unlock() + + state := rl.getOrCreateState(tokenKey) + state.FailCount++ + state.CooldownEnd = time.Now().Add(rl.calculateBackoff(state.FailCount)) +} + +// MarkTokenSuccess 标记 Token 成功 +func (rl *RateLimiter) MarkTokenSuccess(tokenKey string) { + rl.mu.Lock() + defer rl.mu.Unlock() + + state := rl.getOrCreateState(tokenKey) + state.FailCount = 0 + state.CooldownEnd = time.Time{} +} + +// CheckAndMarkSuspended 检测暂停错误并标记 +func (rl *RateLimiter) CheckAndMarkSuspended(tokenKey string, errorMsg string) bool { + suspendKeywords := []string{ + "suspended", + "banned", + "disabled", + "account has been", + "access denied", + "rate limit exceeded", + "too many requests", + "quota exceeded", + } + + lowerMsg := strings.ToLower(errorMsg) + for _, keyword := range suspendKeywords { + if strings.Contains(lowerMsg, keyword) { + rl.mu.Lock() + defer rl.mu.Unlock() + + state := rl.getOrCreateState(tokenKey) + state.IsSuspended = true + state.SuspendedAt = time.Now() + state.SuspendReason = errorMsg + state.CooldownEnd = time.Now().Add(rl.suspendCooldown) + return true + } + } + return false +} + +// IsTokenAvailable 检查 Token 是否可用 +func (rl *RateLimiter) IsTokenAvailable(tokenKey string) bool { + rl.mu.RLock() + defer rl.mu.RUnlock() + + state, exists := rl.states[tokenKey] + if !exists { + return true + } + + now := time.Now() + + // 检查是否被暂停 + if state.IsSuspended { + if now.After(state.SuspendedAt.Add(rl.suspendCooldown)) { + return true + } + return false + } + + // 检查是否在冷却期 + if now.Before(state.CooldownEnd) { + return false + } + + // 检查每日请求限制 + rl.mu.RUnlock() + rl.mu.Lock() + rl.resetDailyIfNeeded(state) + dailyRequests := state.DailyRequests + dailyMax := rl.dailyMaxRequests + rl.mu.Unlock() + rl.mu.RLock() + + if dailyRequests >= dailyMax { + return false + } + + return true +} + +// calculateBackoff 计算指数退避时间 +func (rl *RateLimiter) calculateBackoff(failCount int) time.Duration { + if failCount <= 0 { + return 0 + } + + backoff := float64(rl.backoffBase) * math.Pow(rl.backoffMultiplier, float64(failCount-1)) + + // 添加抖动 + jitter := backoff * rl.jitterPercent * (rl.rng.Float64()*2 - 1) + backoff += jitter + + if time.Duration(backoff) > rl.backoffMax { + return rl.backoffMax + } + return time.Duration(backoff) +} + +// GetTokenState 获取 Token 状态（只读） +func (rl *RateLimiter) GetTokenState(tokenKey string) *TokenState { + rl.mu.RLock() + defer rl.mu.RUnlock() + + state, exists := rl.states[tokenKey] + if !exists { + return nil + } + + // 返回副本以防止外部修改 + stateCopy := *state + return &stateCopy +} + +// ClearTokenState 清除 Token 状态 +func (rl *RateLimiter) ClearTokenState(tokenKey string) { + rl.mu.Lock() + defer rl.mu.Unlock() + delete(rl.states, tokenKey) +} + +// ResetSuspension 重置暂停状态 +func (rl *RateLimiter) ResetSuspension(tokenKey string) { + rl.mu.Lock() + defer rl.mu.Unlock() + + state, exists := rl.states[tokenKey] + if exists { + state.IsSuspended = false + state.SuspendedAt = time.Time{} + state.SuspendReason = "" + state.CooldownEnd = time.Time{} + state.FailCount = 0 + } +} diff --git a/internal/auth/kiro/rate_limiter_singleton.go b/internal/auth/kiro/rate_limiter_singleton.go new file mode 100644 index 00000000..4c02af89 --- /dev/null +++ b/internal/auth/kiro/rate_limiter_singleton.go @@ -0,0 +1,46 @@ +package kiro + +import ( + "sync" + "time" + + log "github.com/sirupsen/logrus" +) + +var ( + globalRateLimiter *RateLimiter + globalRateLimiterOnce sync.Once + + globalCooldownManager *CooldownManager + globalCooldownManagerOnce sync.Once + cooldownStopCh chan struct{} +) + +// GetGlobalRateLimiter returns the singleton RateLimiter instance. +func GetGlobalRateLimiter() *RateLimiter { + globalRateLimiterOnce.Do(func() { + globalRateLimiter = NewRateLimiter() + log.Info("kiro: global RateLimiter initialized") + }) + return globalRateLimiter +} + +// GetGlobalCooldownManager returns the singleton CooldownManager instance. +func GetGlobalCooldownManager() *CooldownManager { + globalCooldownManagerOnce.Do(func() { + globalCooldownManager = NewCooldownManager() + cooldownStopCh = make(chan struct{}) + go globalCooldownManager.StartCleanupRoutine(5*time.Minute, cooldownStopCh) + log.Info("kiro: global CooldownManager initialized with cleanup routine") + }) + return globalCooldownManager +} + +// ShutdownRateLimiters stops the cooldown cleanup routine. +// Should be called during application shutdown. +func ShutdownRateLimiters() { + if cooldownStopCh != nil { + close(cooldownStopCh) + log.Info("kiro: rate limiter cleanup routine stopped") + } +} diff --git a/internal/auth/kiro/rate_limiter_test.go b/internal/auth/kiro/rate_limiter_test.go new file mode 100644 index 00000000..636413dd --- /dev/null +++ b/internal/auth/kiro/rate_limiter_test.go @@ -0,0 +1,304 @@ +package kiro + +import ( + "sync" + "testing" + "time" +) + +func TestNewRateLimiter(t *testing.T) { + rl := NewRateLimiter() + if rl == nil { + t.Fatal("expected non-nil RateLimiter") + } + if rl.states == nil { + t.Error("expected non-nil states map") + } + if rl.minTokenInterval != DefaultMinTokenInterval { + t.Errorf("expected minTokenInterval %v, got %v", DefaultMinTokenInterval, rl.minTokenInterval) + } + if rl.maxTokenInterval != DefaultMaxTokenInterval { + t.Errorf("expected maxTokenInterval %v, got %v", DefaultMaxTokenInterval, rl.maxTokenInterval) + } + if rl.dailyMaxRequests != DefaultDailyMaxRequests { + t.Errorf("expected dailyMaxRequests %d, got %d", DefaultDailyMaxRequests, rl.dailyMaxRequests) + } +} + +func TestNewRateLimiterWithConfig(t *testing.T) { + cfg := RateLimiterConfig{ + MinTokenInterval: 5 * time.Second, + MaxTokenInterval: 15 * time.Second, + DailyMaxRequests: 100, + JitterPercent: 0.2, + BackoffBase: 1 * time.Minute, + BackoffMax: 30 * time.Minute, + BackoffMultiplier: 1.5, + SuspendCooldown: 12 * time.Hour, + } + + rl := NewRateLimiterWithConfig(cfg) + if rl.minTokenInterval != 5*time.Second { + t.Errorf("expected minTokenInterval 5s, got %v", rl.minTokenInterval) + } + if rl.maxTokenInterval != 15*time.Second { + t.Errorf("expected maxTokenInterval 15s, got %v", rl.maxTokenInterval) + } + if rl.dailyMaxRequests != 100 { + t.Errorf("expected dailyMaxRequests 100, got %d", rl.dailyMaxRequests) + } +} + +func TestNewRateLimiterWithConfig_PartialConfig(t *testing.T) { + cfg := RateLimiterConfig{ + MinTokenInterval: 5 * time.Second, + } + + rl := NewRateLimiterWithConfig(cfg) + if rl.minTokenInterval != 5*time.Second { + t.Errorf("expected minTokenInterval 5s, got %v", rl.minTokenInterval) + } + if rl.maxTokenInterval != DefaultMaxTokenInterval { + t.Errorf("expected default maxTokenInterval, got %v", rl.maxTokenInterval) + } +} + +func TestGetTokenState_NonExistent(t *testing.T) { + rl := NewRateLimiter() + state := rl.GetTokenState("nonexistent") + if state != nil { + t.Error("expected nil state for non-existent token") + } +} + +func TestIsTokenAvailable_NewToken(t *testing.T) { + rl := NewRateLimiter() + if !rl.IsTokenAvailable("newtoken") { + t.Error("expected new token to be available") + } +} + +func TestMarkTokenFailed(t *testing.T) { + rl := NewRateLimiter() + rl.MarkTokenFailed("token1") + + state := rl.GetTokenState("token1") + if state == nil { + t.Fatal("expected non-nil state") + } + if state.FailCount != 1 { + t.Errorf("expected FailCount 1, got %d", state.FailCount) + } + if state.CooldownEnd.IsZero() { + t.Error("expected non-zero CooldownEnd") + } +} + +func TestMarkTokenSuccess(t *testing.T) { + rl := NewRateLimiter() + rl.MarkTokenFailed("token1") + rl.MarkTokenFailed("token1") + rl.MarkTokenSuccess("token1") + + state := rl.GetTokenState("token1") + if state == nil { + t.Fatal("expected non-nil state") + } + if state.FailCount != 0 { + t.Errorf("expected FailCount 0, got %d", state.FailCount) + } + if !state.CooldownEnd.IsZero() { + t.Error("expected zero CooldownEnd after success") + } +} + +func TestCheckAndMarkSuspended_Suspended(t *testing.T) { + rl := NewRateLimiter() + + testCases := []string{ + "Account has been suspended", + "You are banned from this service", + "Account disabled", + "Access denied permanently", + "Rate limit exceeded", + "Too many requests", + "Quota exceeded for today", + } + + for i, msg := range testCases { + tokenKey := "token" + string(rune('a'+i)) + if !rl.CheckAndMarkSuspended(tokenKey, msg) { + t.Errorf("expected suspension detected for: %s", msg) + } + state := rl.GetTokenState(tokenKey) + if !state.IsSuspended { + t.Errorf("expected IsSuspended true for: %s", msg) + } + } +} + +func TestCheckAndMarkSuspended_NotSuspended(t *testing.T) { + rl := NewRateLimiter() + + normalErrors := []string{ + "connection timeout", + "internal server error", + "bad request", + "invalid token format", + } + + for i, msg := range normalErrors { + tokenKey := "token" + string(rune('a'+i)) + if rl.CheckAndMarkSuspended(tokenKey, msg) { + t.Errorf("unexpected suspension for: %s", msg) + } + } +} + +func TestIsTokenAvailable_Suspended(t *testing.T) { + rl := NewRateLimiter() + rl.CheckAndMarkSuspended("token1", "Account suspended") + + if rl.IsTokenAvailable("token1") { + t.Error("expected suspended token to be unavailable") + } +} + +func TestClearTokenState(t *testing.T) { + rl := NewRateLimiter() + rl.MarkTokenFailed("token1") + rl.ClearTokenState("token1") + + state := rl.GetTokenState("token1") + if state != nil { + t.Error("expected nil state after clear") + } +} + +func TestResetSuspension(t *testing.T) { + rl := NewRateLimiter() + rl.CheckAndMarkSuspended("token1", "Account suspended") + rl.ResetSuspension("token1") + + state := rl.GetTokenState("token1") + if state.IsSuspended { + t.Error("expected IsSuspended false after reset") + } + if state.FailCount != 0 { + t.Errorf("expected FailCount 0, got %d", state.FailCount) + } +} + +func TestResetSuspension_NonExistent(t *testing.T) { + rl := NewRateLimiter() + rl.ResetSuspension("nonexistent") +} + +func TestCalculateBackoff_ZeroFailCount(t *testing.T) { + rl := NewRateLimiter() + backoff := rl.calculateBackoff(0) + if backoff != 0 { + t.Errorf("expected 0 backoff for 0 fails, got %v", backoff) + } +} + +func TestCalculateBackoff_Exponential(t *testing.T) { + cfg := RateLimiterConfig{ + BackoffBase: 1 * time.Minute, + BackoffMax: 60 * time.Minute, + BackoffMultiplier: 2.0, + JitterPercent: 0.3, + } + rl := NewRateLimiterWithConfig(cfg) + + backoff1 := rl.calculateBackoff(1) + if backoff1 < 40*time.Second || backoff1 > 80*time.Second { + t.Errorf("expected ~1min (with jitter) for fail 1, got %v", backoff1) + } + + backoff2 := rl.calculateBackoff(2) + if backoff2 < 80*time.Second || backoff2 > 160*time.Second { + t.Errorf("expected ~2min (with jitter) for fail 2, got %v", backoff2) + } +} + +func TestCalculateBackoff_MaxCap(t *testing.T) { + cfg := RateLimiterConfig{ + BackoffBase: 1 * time.Minute, + BackoffMax: 10 * time.Minute, + BackoffMultiplier: 2.0, + JitterPercent: 0, + } + rl := NewRateLimiterWithConfig(cfg) + + backoff := rl.calculateBackoff(10) + if backoff > 10*time.Minute { + t.Errorf("expected backoff capped at 10min, got %v", backoff) + } +} + +func TestGetTokenState_ReturnsCopy(t *testing.T) { + rl := NewRateLimiter() + rl.MarkTokenFailed("token1") + + state1 := rl.GetTokenState("token1") + state1.FailCount = 999 + + state2 := rl.GetTokenState("token1") + if state2.FailCount == 999 { + t.Error("GetTokenState should return a copy") + } +} + +func TestRateLimiter_ConcurrentAccess(t *testing.T) { + rl := NewRateLimiter() + const numGoroutines = 50 + const numOperations = 50 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + tokenKey := "token" + string(rune('a'+id%10)) + for j := 0; j < numOperations; j++ { + switch j % 6 { + case 0: + rl.IsTokenAvailable(tokenKey) + case 1: + rl.MarkTokenFailed(tokenKey) + case 2: + rl.MarkTokenSuccess(tokenKey) + case 3: + rl.GetTokenState(tokenKey) + case 4: + rl.CheckAndMarkSuspended(tokenKey, "test error") + case 5: + rl.ResetSuspension(tokenKey) + } + } + }(i) + } + + wg.Wait() +} + +func TestCalculateInterval_WithinRange(t *testing.T) { + cfg := RateLimiterConfig{ + MinTokenInterval: 10 * time.Second, + MaxTokenInterval: 30 * time.Second, + JitterPercent: 0.3, + } + rl := NewRateLimiterWithConfig(cfg) + + minAllowed := 7 * time.Second + maxAllowed := 40 * time.Second + + for i := 0; i < 100; i++ { + interval := rl.calculateInterval() + if interval < minAllowed || interval > maxAllowed { + t.Errorf("interval %v outside expected range [%v, %v]", interval, minAllowed, maxAllowed) + } + } +} diff --git a/internal/auth/kiro/social_auth.go b/internal/auth/kiro/social_auth.go index 2ac29bf8..277b83db 100644 --- a/internal/auth/kiro/social_auth.go +++ b/internal/auth/kiro/social_auth.go @@ -9,7 +9,9 @@ import ( "encoding/base64" "encoding/json" "fmt" + "html" "io" + "net" "net/http" "net/url" "os" @@ -31,6 +33,9 @@ const ( // OAuth timeout socialAuthTimeout = 10 * time.Minute + + // Default callback port for social auth HTTP server + socialAuthCallbackPort = 9876 ) // SocialProvider represents the social login provider. @@ -67,6 +72,13 @@ type RefreshTokenRequest struct { RefreshToken string `json:"refreshToken"` } +// WebCallbackResult contains the OAuth callback result from HTTP server. +type WebCallbackResult struct { + Code string + State string + Error string +} + // SocialAuthClient handles social authentication with Kiro. type SocialAuthClient struct { httpClient *http.Client @@ -87,6 +99,83 @@ func NewSocialAuthClient(cfg *config.Config) *SocialAuthClient { } } +// startWebCallbackServer starts a local HTTP server to receive the OAuth callback. +// This is used instead of the kiro:// protocol handler to avoid redirect_mismatch errors. +func (c *SocialAuthClient) startWebCallbackServer(ctx context.Context, expectedState string) (string, <-chan WebCallbackResult, error) { + // Try to find an available port - use localhost like Kiro does + listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", socialAuthCallbackPort)) + if err != nil { + // Try with dynamic port (RFC 8252 allows dynamic ports for native apps) + log.Warnf("kiro social auth: default port %d is busy, falling back to dynamic port", socialAuthCallbackPort) + listener, err = net.Listen("tcp", "localhost:0") + if err != nil { + return "", nil, fmt.Errorf("failed to start callback server: %w", err) + } + } + + port := listener.Addr().(*net.TCPAddr).Port + // Use http scheme for local callback server + redirectURI := fmt.Sprintf("http://localhost:%d/oauth/callback", port) + resultChan := make(chan WebCallbackResult, 1) + + server := &http.Server{ + ReadHeaderTimeout: 10 * time.Second, + } + + mux := http.NewServeMux() + mux.HandleFunc("/oauth/callback", func(w http.ResponseWriter, r *http.Request) { + code := r.URL.Query().Get("code") + state := r.URL.Query().Get("state") + errParam := r.URL.Query().Get("error") + + if errParam != "" { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.WriteHeader(http.StatusBadRequest) + fmt.Fprintf(w, ` +Login Failed +

Login Failed

%s

You can close this window.

`, html.EscapeString(errParam)) + resultChan <- WebCallbackResult{Error: errParam} + return + } + + if state != expectedState { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.WriteHeader(http.StatusBadRequest) + fmt.Fprint(w, ` +Login Failed +

Login Failed

Invalid state parameter

You can close this window.

`) + resultChan <- WebCallbackResult{Error: "state mismatch"} + return + } + + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, ` +Login Successful +

Login Successful!

You can close this window and return to the terminal.

+`) + resultChan <- WebCallbackResult{Code: code, State: state} + }) + + server.Handler = mux + + go func() { + if err := server.Serve(listener); err != nil && err != http.ErrServerClosed { + log.Debugf("kiro social auth callback server error: %v", err) + } + }() + + go func() { + select { + case <-ctx.Done(): + case <-time.After(socialAuthTimeout): + case <-resultChan: + } + _ = server.Shutdown(context.Background()) + }() + + return redirectURI, resultChan, nil +} + // generatePKCE generates PKCE code verifier and challenge. func generatePKCE() (verifier, challenge string, err error) { // Generate 32 bytes of random data for verifier @@ -217,10 +306,12 @@ func (c *SocialAuthClient) RefreshSocialToken(ctx context.Context, refreshToken ExpiresAt: expiresAt.Format(time.RFC3339), AuthMethod: "social", Provider: "", // Caller should preserve original provider + Region: "us-east-1", }, nil } -// LoginWithSocial performs OAuth login with Google. +// LoginWithSocial performs OAuth login with Google or GitHub. +// Uses local HTTP callback server instead of custom protocol handler to avoid redirect_mismatch errors. func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialProvider) (*KiroTokenData, error) { providerName := string(provider) @@ -228,28 +319,10 @@ func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialP fmt.Printf("║ Kiro Authentication (%s) ║\n", providerName) fmt.Println("╚══════════════════════════════════════════════════════════╝") - // Step 1: Setup protocol handler + // Step 1: Start local HTTP callback server (instead of kiro:// protocol handler) + // This avoids redirect_mismatch errors with AWS Cognito fmt.Println("\nSetting up authentication...") - // Start the local callback server - handlerPort, err := c.protocolHandler.Start(ctx) - if err != nil { - return nil, fmt.Errorf("failed to start callback server: %w", err) - } - defer c.protocolHandler.Stop() - - // Ensure protocol handler is installed and set as default - if err := SetupProtocolHandlerIfNeeded(handlerPort); err != nil { - fmt.Println("\n⚠ Protocol handler setup failed. Trying alternative method...") - fmt.Println(" If you see a browser 'Open with' dialog, select your default browser.") - fmt.Println(" For manual setup instructions, run: cliproxy kiro --help-protocol") - log.Debugf("kiro: protocol handler setup error: %v", err) - // Continue anyway - user might have set it up manually or select browser manually - } else { - // Force set our handler as default (prevents "Open with" dialog) - forceDefaultProtocolHandler() - } - // Step 2: Generate PKCE codes codeVerifier, codeChallenge, err := generatePKCE() if err != nil { @@ -262,8 +335,15 @@ func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialP return nil, fmt.Errorf("failed to generate state: %w", err) } - // Step 4: Build the login URL (Kiro uses GET request with query params) - authURL := c.buildLoginURL(providerName, KiroRedirectURI, codeChallenge, state) + // Step 4: Start local HTTP callback server + redirectURI, resultChan, err := c.startWebCallbackServer(ctx, state) + if err != nil { + return nil, fmt.Errorf("failed to start callback server: %w", err) + } + log.Debugf("kiro social auth: callback server started at %s", redirectURI) + + // Step 5: Build the login URL using HTTP redirect URI + authURL := c.buildLoginURL(providerName, redirectURI, codeChallenge, state) // Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito) // Incognito mode enables multi-account support by bypassing cached sessions @@ -279,7 +359,7 @@ func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialP log.Debug("kiro: using incognito mode for multi-account support (default)") } - // Step 5: Open browser for user authentication + // Step 6: Open browser for user authentication fmt.Println("\n════════════════════════════════════════════════════════════") fmt.Printf(" Opening browser for %s authentication...\n", providerName) fmt.Println("════════════════════════════════════════════════════════════") @@ -295,80 +375,78 @@ func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialP fmt.Println("\n Waiting for authentication callback...") - // Step 6: Wait for callback - callback, err := c.protocolHandler.WaitForCallback(ctx) - if err != nil { - return nil, fmt.Errorf("failed to receive callback: %w", err) - } - - if callback.Error != "" { - return nil, fmt.Errorf("authentication error: %s", callback.Error) - } - - if callback.State != state { - // Log state values for debugging, but don't expose in user-facing error - log.Debugf("kiro: OAuth state mismatch - expected %s, got %s", state, callback.State) - return nil, fmt.Errorf("OAuth state validation failed - please try again") - } - - if callback.Code == "" { - return nil, fmt.Errorf("no authorization code received") - } - - fmt.Println("\n✓ Authorization received!") - - // Step 7: Exchange code for tokens - fmt.Println("Exchanging code for tokens...") - - tokenReq := &CreateTokenRequest{ - Code: callback.Code, - CodeVerifier: codeVerifier, - RedirectURI: KiroRedirectURI, - } - - tokenResp, err := c.CreateToken(ctx, tokenReq) - if err != nil { - return nil, fmt.Errorf("failed to exchange code for tokens: %w", err) - } - - fmt.Println("\n✓ Authentication successful!") - - // Close the browser window - if err := browser.CloseBrowser(); err != nil { - log.Debugf("Failed to close browser: %v", err) - } - - // Validate ExpiresIn - use default 1 hour if invalid - expiresIn := tokenResp.ExpiresIn - if expiresIn <= 0 { - expiresIn = 3600 - } - expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second) - - // Try to extract email from JWT access token first - email := ExtractEmailFromJWT(tokenResp.AccessToken) - - // If no email in JWT, ask user for account label (only in interactive mode) - if email == "" && isInteractiveTerminal() { - fmt.Print("\n Enter account label for file naming (optional, press Enter to skip): ") - reader := bufio.NewReader(os.Stdin) - var err error - email, err = reader.ReadString('\n') - if err != nil { - log.Debugf("Failed to read account label: %v", err) + // Step 7: Wait for callback from HTTP server + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(socialAuthTimeout): + return nil, fmt.Errorf("authentication timed out") + case callback := <-resultChan: + if callback.Error != "" { + return nil, fmt.Errorf("authentication error: %s", callback.Error) } - email = strings.TrimSpace(email) - } - return &KiroTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - ProfileArn: tokenResp.ProfileArn, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: "social", - Provider: providerName, - Email: email, // JWT email or user-provided label - }, nil + // State is already validated by the callback server + if callback.Code == "" { + return nil, fmt.Errorf("no authorization code received") + } + + fmt.Println("\n✓ Authorization received!") + + // Step 8: Exchange code for tokens + fmt.Println("Exchanging code for tokens...") + + tokenReq := &CreateTokenRequest{ + Code: callback.Code, + CodeVerifier: codeVerifier, + RedirectURI: redirectURI, // Use HTTP redirect URI, not kiro:// protocol + } + + tokenResp, err := c.CreateToken(ctx, tokenReq) + if err != nil { + return nil, fmt.Errorf("failed to exchange code for tokens: %w", err) + } + + fmt.Println("\n✓ Authentication successful!") + + // Close the browser window + if err := browser.CloseBrowser(); err != nil { + log.Debugf("Failed to close browser: %v", err) + } + + // Validate ExpiresIn - use default 1 hour if invalid + expiresIn := tokenResp.ExpiresIn + if expiresIn <= 0 { + expiresIn = 3600 + } + expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second) + + // Try to extract email from JWT access token first + email := ExtractEmailFromJWT(tokenResp.AccessToken) + + // If no email in JWT, ask user for account label (only in interactive mode) + if email == "" && isInteractiveTerminal() { + fmt.Print("\n Enter account label for file naming (optional, press Enter to skip): ") + reader := bufio.NewReader(os.Stdin) + var err error + email, err = reader.ReadString('\n') + if err != nil { + log.Debugf("Failed to read account label: %v", err) + } + email = strings.TrimSpace(email) + } + + return &KiroTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: tokenResp.ProfileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: "social", + Provider: providerName, + Email: email, // JWT email or user-provided label + Region: "us-east-1", + }, nil + } } // LoginWithGoogle performs OAuth login with Google. diff --git a/internal/auth/kiro/sso_oidc.go b/internal/auth/kiro/sso_oidc.go index ab44e55f..ba15dac9 100644 --- a/internal/auth/kiro/sso_oidc.go +++ b/internal/auth/kiro/sso_oidc.go @@ -735,6 +735,7 @@ func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret Provider: "AWS", ClientID: clientID, ClientSecret: clientSecret, + Region: defaultIDCRegion, }, nil } @@ -850,16 +851,17 @@ func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, ClientID: regResp.ClientID, ClientSecret: regResp.ClientSecret, Email: email, + Region: defaultIDCRegion, }, nil - } - } + } + } - // Close browser on timeout for better UX - if err := browser.CloseBrowser(); err != nil { - log.Debugf("Failed to close browser on timeout: %v", err) - } - return nil, fmt.Errorf("authorization timed out") -} + // Close browser on timeout for better UX + if err := browser.CloseBrowser(); err != nil { + log.Debugf("Failed to close browser on timeout: %v", err) + } + return nil, fmt.Errorf("authorization timed out") + } // FetchUserEmail retrieves the user's email from AWS SSO OIDC userinfo endpoint. // Falls back to JWT parsing if userinfo fails. @@ -1366,6 +1368,7 @@ func (c *SSOOIDCClient) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTo ClientID: regResp.ClientID, ClientSecret: regResp.ClientSecret, Email: email, + Region: defaultIDCRegion, }, nil } } diff --git a/internal/auth/kiro/sso_oidc.go.bak b/internal/auth/kiro/sso_oidc.go.bak new file mode 100644 index 00000000..ab44e55f --- /dev/null +++ b/internal/auth/kiro/sso_oidc.go.bak @@ -0,0 +1,1371 @@ +// Package kiro provides AWS SSO OIDC authentication for Kiro. +package kiro + +import ( + "bufio" + "context" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "html" + "io" + "net" + "net/http" + "os" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/browser" + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + "github.com/router-for-me/CLIProxyAPI/v6/internal/util" + log "github.com/sirupsen/logrus" +) + +const ( + // AWS SSO OIDC endpoints + ssoOIDCEndpoint = "https://oidc.us-east-1.amazonaws.com" + + // Kiro's start URL for Builder ID + builderIDStartURL = "https://view.awsapps.com/start" + + // Default region for IDC + defaultIDCRegion = "us-east-1" + + // Polling interval + pollInterval = 5 * time.Second + + // Authorization code flow callback + authCodeCallbackPath = "/oauth/callback" + authCodeCallbackPort = 19877 + + // User-Agent to match official Kiro IDE + kiroUserAgent = "KiroIDE" + + // IDC token refresh headers (matching Kiro IDE behavior) + idcAmzUserAgent = "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE" +) + +// Sentinel errors for OIDC token polling +var ( + ErrAuthorizationPending = errors.New("authorization_pending") + ErrSlowDown = errors.New("slow_down") +) + +// SSOOIDCClient handles AWS SSO OIDC authentication. +type SSOOIDCClient struct { + httpClient *http.Client + cfg *config.Config +} + +// NewSSOOIDCClient creates a new SSO OIDC client. +func NewSSOOIDCClient(cfg *config.Config) *SSOOIDCClient { + client := &http.Client{Timeout: 30 * time.Second} + if cfg != nil { + client = util.SetProxy(&cfg.SDKConfig, client) + } + return &SSOOIDCClient{ + httpClient: client, + cfg: cfg, + } +} + +// RegisterClientResponse from AWS SSO OIDC. +type RegisterClientResponse struct { + ClientID string `json:"clientId"` + ClientSecret string `json:"clientSecret"` + ClientIDIssuedAt int64 `json:"clientIdIssuedAt"` + ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"` +} + +// StartDeviceAuthResponse from AWS SSO OIDC. +type StartDeviceAuthResponse struct { + DeviceCode string `json:"deviceCode"` + UserCode string `json:"userCode"` + VerificationURI string `json:"verificationUri"` + VerificationURIComplete string `json:"verificationUriComplete"` + ExpiresIn int `json:"expiresIn"` + Interval int `json:"interval"` +} + +// CreateTokenResponse from AWS SSO OIDC. +type CreateTokenResponse struct { + AccessToken string `json:"accessToken"` + TokenType string `json:"tokenType"` + ExpiresIn int `json:"expiresIn"` + RefreshToken string `json:"refreshToken"` +} + +// getOIDCEndpoint returns the OIDC endpoint for the given region. +func getOIDCEndpoint(region string) string { + if region == "" { + region = defaultIDCRegion + } + return fmt.Sprintf("https://oidc.%s.amazonaws.com", region) +} + +// promptInput prompts the user for input with an optional default value. +func promptInput(prompt, defaultValue string) string { + reader := bufio.NewReader(os.Stdin) + if defaultValue != "" { + fmt.Printf("%s [%s]: ", prompt, defaultValue) + } else { + fmt.Printf("%s: ", prompt) + } + input, err := reader.ReadString('\n') + if err != nil { + log.Warnf("Error reading input: %v", err) + return defaultValue + } + input = strings.TrimSpace(input) + if input == "" { + return defaultValue + } + return input +} + +// promptSelect prompts the user to select from options using number input. +func promptSelect(prompt string, options []string) int { + reader := bufio.NewReader(os.Stdin) + + for { + fmt.Println(prompt) + for i, opt := range options { + fmt.Printf(" %d) %s\n", i+1, opt) + } + fmt.Printf("Enter selection (1-%d): ", len(options)) + + input, err := reader.ReadString('\n') + if err != nil { + log.Warnf("Error reading input: %v", err) + return 0 // Default to first option on error + } + input = strings.TrimSpace(input) + + // Parse the selection + var selection int + if _, err := fmt.Sscanf(input, "%d", &selection); err != nil || selection < 1 || selection > len(options) { + fmt.Printf("Invalid selection '%s'. Please enter a number between 1 and %d.\n\n", input, len(options)) + continue + } + return selection - 1 + } +} + +// RegisterClientWithRegion registers a new OIDC client with AWS using a specific region. +func (c *SSOOIDCClient) RegisterClientWithRegion(ctx context.Context, region string) (*RegisterClientResponse, error) { + endpoint := getOIDCEndpoint(region) + + payload := map[string]interface{}{ + "clientName": "Kiro IDE", + "clientType": "public", + "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"}, + "grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"}, + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/client/register", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode) + } + + var result RegisterClientResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// StartDeviceAuthorizationWithIDC starts the device authorization flow for IDC. +func (c *SSOOIDCClient) StartDeviceAuthorizationWithIDC(ctx context.Context, clientID, clientSecret, startURL, region string) (*StartDeviceAuthResponse, error) { + endpoint := getOIDCEndpoint(region) + + payload := map[string]string{ + "clientId": clientID, + "clientSecret": clientSecret, + "startUrl": startURL, + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/device_authorization", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode) + } + + var result StartDeviceAuthResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// CreateTokenWithRegion polls for the access token after user authorization using a specific region. +func (c *SSOOIDCClient) CreateTokenWithRegion(ctx context.Context, clientID, clientSecret, deviceCode, region string) (*CreateTokenResponse, error) { + endpoint := getOIDCEndpoint(region) + + payload := map[string]string{ + "clientId": clientID, + "clientSecret": clientSecret, + "deviceCode": deviceCode, + "grantType": "urn:ietf:params:oauth:grant-type:device_code", + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + // Check for pending authorization + if resp.StatusCode == http.StatusBadRequest { + var errResp struct { + Error string `json:"error"` + } + if json.Unmarshal(respBody, &errResp) == nil { + if errResp.Error == "authorization_pending" { + return nil, ErrAuthorizationPending + } + if errResp.Error == "slow_down" { + return nil, ErrSlowDown + } + } + log.Debugf("create token failed: %s", string(respBody)) + return nil, fmt.Errorf("create token failed") + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode) + } + + var result CreateTokenResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// RefreshTokenWithRegion refreshes an access token using the refresh token with a specific region. +func (c *SSOOIDCClient) RefreshTokenWithRegion(ctx context.Context, clientID, clientSecret, refreshToken, region, startURL string) (*KiroTokenData, error) { + endpoint := getOIDCEndpoint(region) + + payload := map[string]string{ + "clientId": clientID, + "clientSecret": clientSecret, + "refreshToken": refreshToken, + "grantType": "refresh_token", + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + + // Set headers matching kiro2api's IDC token refresh + // These headers are required for successful IDC token refresh + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region)) + req.Header.Set("Connection", "keep-alive") + req.Header.Set("x-amz-user-agent", idcAmzUserAgent) + req.Header.Set("Accept", "*/*") + req.Header.Set("Accept-Language", "*") + req.Header.Set("sec-fetch-mode", "cors") + req.Header.Set("User-Agent", "node") + req.Header.Set("Accept-Encoding", "br, gzip, deflate") + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Warnf("IDC token refresh failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode) + } + + var result CreateTokenResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second) + + return &KiroTokenData{ + AccessToken: result.AccessToken, + RefreshToken: result.RefreshToken, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: "idc", + Provider: "AWS", + ClientID: clientID, + ClientSecret: clientSecret, + StartURL: startURL, + Region: region, + }, nil +} + +// LoginWithIDC performs the full device code flow for AWS Identity Center (IDC). +func (c *SSOOIDCClient) LoginWithIDC(ctx context.Context, startURL, region string) (*KiroTokenData, error) { + fmt.Println("\n╔══════════════════════════════════════════════════════════╗") + fmt.Println("║ Kiro Authentication (AWS Identity Center) ║") + fmt.Println("╚══════════════════════════════════════════════════════════╝") + + // Step 1: Register client with the specified region + fmt.Println("\nRegistering client...") + regResp, err := c.RegisterClientWithRegion(ctx, region) + if err != nil { + return nil, fmt.Errorf("failed to register client: %w", err) + } + log.Debugf("Client registered: %s", regResp.ClientID) + + // Step 2: Start device authorization with IDC start URL + fmt.Println("Starting device authorization...") + authResp, err := c.StartDeviceAuthorizationWithIDC(ctx, regResp.ClientID, regResp.ClientSecret, startURL, region) + if err != nil { + return nil, fmt.Errorf("failed to start device auth: %w", err) + } + + // Step 3: Show user the verification URL + fmt.Printf("\n") + fmt.Println("════════════════════════════════════════════════════════════") + fmt.Printf(" Confirm the following code in the browser:\n") + fmt.Printf(" Code: %s\n", authResp.UserCode) + fmt.Println("════════════════════════════════════════════════════════════") + fmt.Printf("\n Open this URL: %s\n\n", authResp.VerificationURIComplete) + + // Set incognito mode based on config + if c.cfg != nil { + browser.SetIncognitoMode(c.cfg.IncognitoBrowser) + if !c.cfg.IncognitoBrowser { + log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.") + } else { + log.Debug("kiro: using incognito mode for multi-account support") + } + } else { + browser.SetIncognitoMode(true) + log.Debug("kiro: using incognito mode for multi-account support (default)") + } + + // Open browser + if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil { + log.Warnf("Could not open browser automatically: %v", err) + fmt.Println(" Please open the URL manually in your browser.") + } else { + fmt.Println(" (Browser opened automatically)") + } + + // Step 4: Poll for token + fmt.Println("Waiting for authorization...") + + interval := pollInterval + if authResp.Interval > 0 { + interval = time.Duration(authResp.Interval) * time.Second + } + + deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second) + + for time.Now().Before(deadline) { + select { + case <-ctx.Done(): + browser.CloseBrowser() + return nil, ctx.Err() + case <-time.After(interval): + tokenResp, err := c.CreateTokenWithRegion(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode, region) + if err != nil { + if errors.Is(err, ErrAuthorizationPending) { + fmt.Print(".") + continue + } + if errors.Is(err, ErrSlowDown) { + interval += 5 * time.Second + continue + } + browser.CloseBrowser() + return nil, fmt.Errorf("token creation failed: %w", err) + } + + fmt.Println("\n\n✓ Authorization successful!") + + // Close the browser window + if err := browser.CloseBrowser(); err != nil { + log.Debugf("Failed to close browser: %v", err) + } + + // Step 5: Get profile ARN from CodeWhisperer API + fmt.Println("Fetching profile information...") + profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken) + + // Fetch user email + email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken) + if email != "" { + fmt.Printf(" Logged in as: %s\n", email) + } + + expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) + + return &KiroTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: profileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: "idc", + Provider: "AWS", + ClientID: regResp.ClientID, + ClientSecret: regResp.ClientSecret, + Email: email, + StartURL: startURL, + Region: region, + }, nil + } + } + + // Close browser on timeout + if err := browser.CloseBrowser(); err != nil { + log.Debugf("Failed to close browser on timeout: %v", err) + } + return nil, fmt.Errorf("authorization timed out") +} + +// LoginWithMethodSelection prompts the user to select between Builder ID and IDC, then performs the login. +func (c *SSOOIDCClient) LoginWithMethodSelection(ctx context.Context) (*KiroTokenData, error) { + fmt.Println("\n╔══════════════════════════════════════════════════════════╗") + fmt.Println("║ Kiro Authentication (AWS) ║") + fmt.Println("╚══════════════════════════════════════════════════════════╝") + + // Prompt for login method + options := []string{ + "Use with Builder ID (personal AWS account)", + "Use with IDC Account (organization SSO)", + } + selection := promptSelect("\n? Select login method:", options) + + if selection == 0 { + // Builder ID flow - use existing implementation + return c.LoginWithBuilderID(ctx) + } + + // IDC flow - prompt for start URL and region + fmt.Println() + startURL := promptInput("? Enter Start URL", "") + if startURL == "" { + return nil, fmt.Errorf("start URL is required for IDC login") + } + + region := promptInput("? Enter Region", defaultIDCRegion) + + return c.LoginWithIDC(ctx, startURL, region) +} + +// RegisterClient registers a new OIDC client with AWS. +func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResponse, error) { + payload := map[string]interface{}{ + "clientName": "Kiro IDE", + "clientType": "public", + "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"}, + "grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"}, + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode) + } + + var result RegisterClientResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// StartDeviceAuthorization starts the device authorization flow. +func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID, clientSecret string) (*StartDeviceAuthResponse, error) { + payload := map[string]string{ + "clientId": clientID, + "clientSecret": clientSecret, + "startUrl": builderIDStartURL, + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/device_authorization", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode) + } + + var result StartDeviceAuthResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// CreateToken polls for the access token after user authorization. +func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret, deviceCode string) (*CreateTokenResponse, error) { + payload := map[string]string{ + "clientId": clientID, + "clientSecret": clientSecret, + "deviceCode": deviceCode, + "grantType": "urn:ietf:params:oauth:grant-type:device_code", + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + // Check for pending authorization + if resp.StatusCode == http.StatusBadRequest { + var errResp struct { + Error string `json:"error"` + } + if json.Unmarshal(respBody, &errResp) == nil { + if errResp.Error == "authorization_pending" { + return nil, ErrAuthorizationPending + } + if errResp.Error == "slow_down" { + return nil, ErrSlowDown + } + } + log.Debugf("create token failed: %s", string(respBody)) + return nil, fmt.Errorf("create token failed") + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode) + } + + var result CreateTokenResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// RefreshToken refreshes an access token using the refresh token. +func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret, refreshToken string) (*KiroTokenData, error) { + payload := map[string]string{ + "clientId": clientID, + "clientSecret": clientSecret, + "refreshToken": refreshToken, + "grantType": "refresh_token", + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode) + } + + var result CreateTokenResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second) + + return &KiroTokenData{ + AccessToken: result.AccessToken, + RefreshToken: result.RefreshToken, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: "builder-id", + Provider: "AWS", + ClientID: clientID, + ClientSecret: clientSecret, + }, nil +} + +// LoginWithBuilderID performs the full device code flow for AWS Builder ID. +func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) { + fmt.Println("\n╔══════════════════════════════════════════════════════════╗") + fmt.Println("║ Kiro Authentication (AWS Builder ID) ║") + fmt.Println("╚══════════════════════════════════════════════════════════╝") + + // Step 1: Register client + fmt.Println("\nRegistering client...") + regResp, err := c.RegisterClient(ctx) + if err != nil { + return nil, fmt.Errorf("failed to register client: %w", err) + } + log.Debugf("Client registered: %s", regResp.ClientID) + + // Step 2: Start device authorization + fmt.Println("Starting device authorization...") + authResp, err := c.StartDeviceAuthorization(ctx, regResp.ClientID, regResp.ClientSecret) + if err != nil { + return nil, fmt.Errorf("failed to start device auth: %w", err) + } + + // Step 3: Show user the verification URL + fmt.Printf("\n") + fmt.Println("════════════════════════════════════════════════════════════") + fmt.Printf(" Open this URL in your browser:\n") + fmt.Printf(" %s\n", authResp.VerificationURIComplete) + fmt.Println("════════════════════════════════════════════════════════════") + fmt.Printf("\n Or go to: %s\n", authResp.VerificationURI) + fmt.Printf(" And enter code: %s\n\n", authResp.UserCode) + + // Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito) + // Incognito mode enables multi-account support by bypassing cached sessions + if c.cfg != nil { + browser.SetIncognitoMode(c.cfg.IncognitoBrowser) + if !c.cfg.IncognitoBrowser { + log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.") + } else { + log.Debug("kiro: using incognito mode for multi-account support") + } + } else { + browser.SetIncognitoMode(true) // Default to incognito if no config + log.Debug("kiro: using incognito mode for multi-account support (default)") + } + + // Open browser using cross-platform browser package + if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil { + log.Warnf("Could not open browser automatically: %v", err) + fmt.Println(" Please open the URL manually in your browser.") + } else { + fmt.Println(" (Browser opened automatically)") + } + + // Step 4: Poll for token + fmt.Println("Waiting for authorization...") + + interval := pollInterval + if authResp.Interval > 0 { + interval = time.Duration(authResp.Interval) * time.Second + } + + deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second) + + for time.Now().Before(deadline) { + select { + case <-ctx.Done(): + browser.CloseBrowser() // Cleanup on cancel + return nil, ctx.Err() + case <-time.After(interval): + tokenResp, err := c.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode) + if err != nil { + if errors.Is(err, ErrAuthorizationPending) { + fmt.Print(".") + continue + } + if errors.Is(err, ErrSlowDown) { + interval += 5 * time.Second + continue + } + // Close browser on error before returning + browser.CloseBrowser() + return nil, fmt.Errorf("token creation failed: %w", err) + } + + fmt.Println("\n\n✓ Authorization successful!") + + // Close the browser window + if err := browser.CloseBrowser(); err != nil { + log.Debugf("Failed to close browser: %v", err) + } + + // Step 5: Get profile ARN from CodeWhisperer API + fmt.Println("Fetching profile information...") + profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken) + + // Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing) + email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken) + if email != "" { + fmt.Printf(" Logged in as: %s\n", email) + } + + expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) + + return &KiroTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: profileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: "builder-id", + Provider: "AWS", + ClientID: regResp.ClientID, + ClientSecret: regResp.ClientSecret, + Email: email, + }, nil + } + } + + // Close browser on timeout for better UX + if err := browser.CloseBrowser(); err != nil { + log.Debugf("Failed to close browser on timeout: %v", err) + } + return nil, fmt.Errorf("authorization timed out") +} + +// FetchUserEmail retrieves the user's email from AWS SSO OIDC userinfo endpoint. +// Falls back to JWT parsing if userinfo fails. +func (c *SSOOIDCClient) FetchUserEmail(ctx context.Context, accessToken string) string { + // Method 1: Try userinfo endpoint (standard OIDC) + email := c.tryUserInfoEndpoint(ctx, accessToken) + if email != "" { + return email + } + + // Method 2: Fallback to JWT parsing + return ExtractEmailFromJWT(accessToken) +} + +// tryUserInfoEndpoint attempts to get user info from AWS SSO OIDC userinfo endpoint. +func (c *SSOOIDCClient) tryUserInfoEndpoint(ctx context.Context, accessToken string) string { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, ssoOIDCEndpoint+"/userinfo", nil) + if err != nil { + return "" + } + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + log.Debugf("userinfo request failed: %v", err) + return "" + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + respBody, _ := io.ReadAll(resp.Body) + log.Debugf("userinfo endpoint returned status %d: %s", resp.StatusCode, string(respBody)) + return "" + } + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return "" + } + + log.Debugf("userinfo response: %s", string(respBody)) + + var userInfo struct { + Email string `json:"email"` + Sub string `json:"sub"` + PreferredUsername string `json:"preferred_username"` + Name string `json:"name"` + } + + if err := json.Unmarshal(respBody, &userInfo); err != nil { + return "" + } + + if userInfo.Email != "" { + return userInfo.Email + } + if userInfo.PreferredUsername != "" && strings.Contains(userInfo.PreferredUsername, "@") { + return userInfo.PreferredUsername + } + return "" +} + +// fetchProfileArn retrieves the profile ARN from CodeWhisperer API. +// This is needed for file naming since AWS SSO OIDC doesn't return profile info. +func (c *SSOOIDCClient) fetchProfileArn(ctx context.Context, accessToken string) string { + // Try ListProfiles API first + profileArn := c.tryListProfiles(ctx, accessToken) + if profileArn != "" { + return profileArn + } + + // Fallback: Try ListAvailableCustomizations + return c.tryListCustomizations(ctx, accessToken) +} + +func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string) string { + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + } + + body, err := json.Marshal(payload) + if err != nil { + return "" + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body))) + if err != nil { + return "" + } + + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListProfiles") + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + return "" + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK { + log.Debugf("ListProfiles failed (status %d): %s", resp.StatusCode, string(respBody)) + return "" + } + + log.Debugf("ListProfiles response: %s", string(respBody)) + + var result struct { + Profiles []struct { + Arn string `json:"arn"` + } `json:"profiles"` + ProfileArn string `json:"profileArn"` + } + + if err := json.Unmarshal(respBody, &result); err != nil { + return "" + } + + if result.ProfileArn != "" { + return result.ProfileArn + } + + if len(result.Profiles) > 0 { + return result.Profiles[0].Arn + } + + return "" +} + +func (c *SSOOIDCClient) tryListCustomizations(ctx context.Context, accessToken string) string { + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + } + + body, err := json.Marshal(payload) + if err != nil { + return "" + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body))) + if err != nil { + return "" + } + + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListAvailableCustomizations") + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + return "" + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK { + log.Debugf("ListAvailableCustomizations failed (status %d): %s", resp.StatusCode, string(respBody)) + return "" + } + + log.Debugf("ListAvailableCustomizations response: %s", string(respBody)) + + var result struct { + Customizations []struct { + Arn string `json:"arn"` + } `json:"customizations"` + ProfileArn string `json:"profileArn"` + } + + if err := json.Unmarshal(respBody, &result); err != nil { + return "" + } + + if result.ProfileArn != "" { + return result.ProfileArn + } + + if len(result.Customizations) > 0 { + return result.Customizations[0].Arn + } + + return "" +} + +// RegisterClientForAuthCode registers a new OIDC client for authorization code flow. +func (c *SSOOIDCClient) RegisterClientForAuthCode(ctx context.Context, redirectURI string) (*RegisterClientResponse, error) { + payload := map[string]interface{}{ + "clientName": "Kiro IDE", + "clientType": "public", + "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"}, + "grantTypes": []string{"authorization_code", "refresh_token"}, + "redirectUris": []string{redirectURI}, + "issuerUrl": builderIDStartURL, + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("register client for auth code failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode) + } + + var result RegisterClientResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// AuthCodeCallbackResult contains the result from authorization code callback. +type AuthCodeCallbackResult struct { + Code string + State string + Error string +} + +// startAuthCodeCallbackServer starts a local HTTP server to receive the authorization code callback. +func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthCodeCallbackResult, error) { + // Try to find an available port + listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", authCodeCallbackPort)) + if err != nil { + // Try with dynamic port + log.Warnf("sso oidc: default port %d is busy, falling back to dynamic port", authCodeCallbackPort) + listener, err = net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return "", nil, fmt.Errorf("failed to start callback server: %w", err) + } + } + + port := listener.Addr().(*net.TCPAddr).Port + redirectURI := fmt.Sprintf("http://127.0.0.1:%d%s", port, authCodeCallbackPath) + resultChan := make(chan AuthCodeCallbackResult, 1) + + server := &http.Server{ + ReadHeaderTimeout: 10 * time.Second, + } + + mux := http.NewServeMux() + mux.HandleFunc(authCodeCallbackPath, func(w http.ResponseWriter, r *http.Request) { + code := r.URL.Query().Get("code") + state := r.URL.Query().Get("state") + errParam := r.URL.Query().Get("error") + + // Send response to browser + w.Header().Set("Content-Type", "text/html; charset=utf-8") + if errParam != "" { + w.WriteHeader(http.StatusBadRequest) + fmt.Fprintf(w, ` +Login Failed +

Login Failed

Error: %s

You can close this window.

`, html.EscapeString(errParam)) + resultChan <- AuthCodeCallbackResult{Error: errParam} + return + } + + if state != expectedState { + w.WriteHeader(http.StatusBadRequest) + fmt.Fprint(w, ` +Login Failed +

Login Failed

Invalid state parameter

You can close this window.

`) + resultChan <- AuthCodeCallbackResult{Error: "state mismatch"} + return + } + + fmt.Fprint(w, ` +Login Successful +

Login Successful!

You can close this window and return to the terminal.

+`) + resultChan <- AuthCodeCallbackResult{Code: code, State: state} + }) + + server.Handler = mux + + go func() { + if err := server.Serve(listener); err != nil && err != http.ErrServerClosed { + log.Debugf("auth code callback server error: %v", err) + } + }() + + go func() { + select { + case <-ctx.Done(): + case <-time.After(10 * time.Minute): + case <-resultChan: + } + _ = server.Shutdown(context.Background()) + }() + + return redirectURI, resultChan, nil +} + +// generatePKCEForAuthCode generates PKCE code verifier and challenge for authorization code flow. +func generatePKCEForAuthCode() (verifier, challenge string, err error) { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return "", "", fmt.Errorf("failed to generate random bytes: %w", err) + } + verifier = base64.RawURLEncoding.EncodeToString(b) + h := sha256.Sum256([]byte(verifier)) + challenge = base64.RawURLEncoding.EncodeToString(h[:]) + return verifier, challenge, nil +} + +// generateStateForAuthCode generates a random state parameter. +func generateStateForAuthCode() (string, error) { + b := make([]byte, 16) + if _, err := rand.Read(b); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(b), nil +} + +// CreateTokenWithAuthCode exchanges authorization code for tokens. +func (c *SSOOIDCClient) CreateTokenWithAuthCode(ctx context.Context, clientID, clientSecret, code, codeVerifier, redirectURI string) (*CreateTokenResponse, error) { + payload := map[string]string{ + "clientId": clientID, + "clientSecret": clientSecret, + "code": code, + "codeVerifier": codeVerifier, + "redirectUri": redirectURI, + "grantType": "authorization_code", + } + + body, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body))) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", kiroUserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + log.Debugf("create token with auth code failed (status %d): %s", resp.StatusCode, string(respBody)) + return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode) + } + + var result CreateTokenResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, err + } + + return &result, nil +} + +// LoginWithBuilderIDAuthCode performs the authorization code flow for AWS Builder ID. +// This provides a better UX than device code flow as it uses automatic browser callback. +func (c *SSOOIDCClient) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTokenData, error) { + fmt.Println("\n╔══════════════════════════════════════════════════════════╗") + fmt.Println("║ Kiro Authentication (AWS Builder ID - Auth Code) ║") + fmt.Println("╚══════════════════════════════════════════════════════════╝") + + // Step 1: Generate PKCE and state + codeVerifier, codeChallenge, err := generatePKCEForAuthCode() + if err != nil { + return nil, fmt.Errorf("failed to generate PKCE: %w", err) + } + + state, err := generateStateForAuthCode() + if err != nil { + return nil, fmt.Errorf("failed to generate state: %w", err) + } + + // Step 2: Start callback server + fmt.Println("\nStarting callback server...") + redirectURI, resultChan, err := c.startAuthCodeCallbackServer(ctx, state) + if err != nil { + return nil, fmt.Errorf("failed to start callback server: %w", err) + } + log.Debugf("Callback server started, redirect URI: %s", redirectURI) + + // Step 3: Register client with auth code grant type + fmt.Println("Registering client...") + regResp, err := c.RegisterClientForAuthCode(ctx, redirectURI) + if err != nil { + return nil, fmt.Errorf("failed to register client: %w", err) + } + log.Debugf("Client registered: %s", regResp.ClientID) + + // Step 4: Build authorization URL + scopes := "codewhisperer:completions,codewhisperer:analysis,codewhisperer:conversations" + authURL := fmt.Sprintf("%s/authorize?response_type=code&client_id=%s&redirect_uri=%s&scopes=%s&state=%s&code_challenge=%s&code_challenge_method=S256", + ssoOIDCEndpoint, + regResp.ClientID, + redirectURI, + scopes, + state, + codeChallenge, + ) + + // Step 5: Open browser + fmt.Println("\n════════════════════════════════════════════════════════════") + fmt.Println(" Opening browser for authentication...") + fmt.Println("════════════════════════════════════════════════════════════") + fmt.Printf("\n URL: %s\n\n", authURL) + + // Set incognito mode + if c.cfg != nil { + browser.SetIncognitoMode(c.cfg.IncognitoBrowser) + } else { + browser.SetIncognitoMode(true) + } + + if err := browser.OpenURL(authURL); err != nil { + log.Warnf("Could not open browser automatically: %v", err) + fmt.Println(" ⚠ Could not open browser automatically.") + fmt.Println(" Please open the URL above in your browser manually.") + } else { + fmt.Println(" (Browser opened automatically)") + } + + fmt.Println("\n Waiting for authorization callback...") + + // Step 6: Wait for callback + select { + case <-ctx.Done(): + browser.CloseBrowser() + return nil, ctx.Err() + case <-time.After(10 * time.Minute): + browser.CloseBrowser() + return nil, fmt.Errorf("authorization timed out") + case result := <-resultChan: + if result.Error != "" { + browser.CloseBrowser() + return nil, fmt.Errorf("authorization failed: %s", result.Error) + } + + fmt.Println("\n✓ Authorization received!") + + // Close browser + if err := browser.CloseBrowser(); err != nil { + log.Debugf("Failed to close browser: %v", err) + } + + // Step 7: Exchange code for tokens + fmt.Println("Exchanging code for tokens...") + tokenResp, err := c.CreateTokenWithAuthCode(ctx, regResp.ClientID, regResp.ClientSecret, result.Code, codeVerifier, redirectURI) + if err != nil { + return nil, fmt.Errorf("failed to exchange code for tokens: %w", err) + } + + fmt.Println("\n✓ Authentication successful!") + + // Step 8: Get profile ARN + fmt.Println("Fetching profile information...") + profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken) + + // Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing) + email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken) + if email != "" { + fmt.Printf(" Logged in as: %s\n", email) + } + + expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) + + return &KiroTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: profileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: "builder-id", + Provider: "AWS", + ClientID: regResp.ClientID, + ClientSecret: regResp.ClientSecret, + Email: email, + }, nil + } +} diff --git a/internal/auth/kiro/token.go b/internal/auth/kiro/token.go index e83b1728..bfbdc795 100644 --- a/internal/auth/kiro/token.go +++ b/internal/auth/kiro/token.go @@ -9,6 +9,8 @@ import ( // KiroTokenStorage holds the persistent token data for Kiro authentication. type KiroTokenStorage struct { + // Type is the provider type for management UI recognition (must be "kiro") + Type string `json:"type"` // AccessToken is the OAuth2 access token for API access AccessToken string `json:"access_token"` // RefreshToken is used to obtain new access tokens @@ -23,6 +25,16 @@ type KiroTokenStorage struct { Provider string `json:"provider"` // LastRefresh is the timestamp of the last token refresh LastRefresh string `json:"last_refresh"` + // ClientID is the OAuth client ID (required for token refresh) + ClientID string `json:"clientId,omitempty"` + // ClientSecret is the OAuth client secret (required for token refresh) + ClientSecret string `json:"clientSecret,omitempty"` + // Region is the AWS region + Region string `json:"region,omitempty"` + // StartURL is the AWS Identity Center start URL (for IDC auth) + StartURL string `json:"startUrl,omitempty"` + // Email is the user's email address + Email string `json:"email,omitempty"` } // SaveTokenToFile persists the token storage to the specified file path. @@ -68,5 +80,10 @@ func (s *KiroTokenStorage) ToTokenData() *KiroTokenData { ExpiresAt: s.ExpiresAt, AuthMethod: s.AuthMethod, Provider: s.Provider, + ClientID: s.ClientID, + ClientSecret: s.ClientSecret, + Region: s.Region, + StartURL: s.StartURL, + Email: s.Email, } } diff --git a/internal/auth/kiro/usage_checker.go b/internal/auth/kiro/usage_checker.go new file mode 100644 index 00000000..94870214 --- /dev/null +++ b/internal/auth/kiro/usage_checker.go @@ -0,0 +1,243 @@ +// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API. +// This file implements usage quota checking and monitoring. +package kiro + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + "github.com/router-for-me/CLIProxyAPI/v6/internal/util" +) + +// UsageQuotaResponse represents the API response structure for usage quota checking. +type UsageQuotaResponse struct { + UsageBreakdownList []UsageBreakdownExtended `json:"usageBreakdownList"` + SubscriptionInfo *SubscriptionInfo `json:"subscriptionInfo,omitempty"` + NextDateReset float64 `json:"nextDateReset,omitempty"` +} + +// UsageBreakdownExtended represents detailed usage information for quota checking. +// Note: UsageBreakdown is already defined in codewhisperer_client.go +type UsageBreakdownExtended struct { + ResourceType string `json:"resourceType"` + UsageLimitWithPrecision float64 `json:"usageLimitWithPrecision"` + CurrentUsageWithPrecision float64 `json:"currentUsageWithPrecision"` + FreeTrialInfo *FreeTrialInfoExtended `json:"freeTrialInfo,omitempty"` +} + +// FreeTrialInfoExtended represents free trial usage information. +type FreeTrialInfoExtended struct { + FreeTrialStatus string `json:"freeTrialStatus"` + UsageLimitWithPrecision float64 `json:"usageLimitWithPrecision"` + CurrentUsageWithPrecision float64 `json:"currentUsageWithPrecision"` +} + +// QuotaStatus represents the quota status for a token. +type QuotaStatus struct { + TotalLimit float64 + CurrentUsage float64 + RemainingQuota float64 + IsExhausted bool + ResourceType string + NextReset time.Time +} + +// UsageChecker provides methods for checking token quota usage. +type UsageChecker struct { + httpClient *http.Client + endpoint string +} + +// NewUsageChecker creates a new UsageChecker instance. +func NewUsageChecker(cfg *config.Config) *UsageChecker { + return &UsageChecker{ + httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 30 * time.Second}), + endpoint: awsKiroEndpoint, + } +} + +// NewUsageCheckerWithClient creates a UsageChecker with a custom HTTP client. +func NewUsageCheckerWithClient(client *http.Client) *UsageChecker { + return &UsageChecker{ + httpClient: client, + endpoint: awsKiroEndpoint, + } +} + +// CheckUsage retrieves usage limits for the given token. +func (c *UsageChecker) CheckUsage(ctx context.Context, tokenData *KiroTokenData) (*UsageQuotaResponse, error) { + if tokenData == nil { + return nil, fmt.Errorf("token data is nil") + } + + if tokenData.AccessToken == "" { + return nil, fmt.Errorf("access token is empty") + } + + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + "profileArn": tokenData.ProfileArn, + "resourceType": "AGENTIC_REQUEST", + } + + jsonBody, err := json.Marshal(payload) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, strings.NewReader(string(jsonBody))) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", targetGetUsage) + req.Header.Set("Authorization", "Bearer "+tokenData.AccessToken) + req.Header.Set("Accept", "application/json") + + 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) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body)) + } + + var result UsageQuotaResponse + if err := json.Unmarshal(body, &result); err != nil { + return nil, fmt.Errorf("failed to parse usage response: %w", err) + } + + return &result, nil +} + +// CheckUsageByAccessToken retrieves usage limits using an access token and profile ARN directly. +func (c *UsageChecker) CheckUsageByAccessToken(ctx context.Context, accessToken, profileArn string) (*UsageQuotaResponse, error) { + tokenData := &KiroTokenData{ + AccessToken: accessToken, + ProfileArn: profileArn, + } + return c.CheckUsage(ctx, tokenData) +} + +// GetRemainingQuota calculates the remaining quota from usage limits. +func GetRemainingQuota(usage *UsageQuotaResponse) float64 { + if usage == nil || len(usage.UsageBreakdownList) == 0 { + return 0 + } + + var totalRemaining float64 + for _, breakdown := range usage.UsageBreakdownList { + remaining := breakdown.UsageLimitWithPrecision - breakdown.CurrentUsageWithPrecision + if remaining > 0 { + totalRemaining += remaining + } + + if breakdown.FreeTrialInfo != nil { + freeRemaining := breakdown.FreeTrialInfo.UsageLimitWithPrecision - breakdown.FreeTrialInfo.CurrentUsageWithPrecision + if freeRemaining > 0 { + totalRemaining += freeRemaining + } + } + } + + return totalRemaining +} + +// IsQuotaExhausted checks if the quota is exhausted based on usage limits. +func IsQuotaExhausted(usage *UsageQuotaResponse) bool { + if usage == nil || len(usage.UsageBreakdownList) == 0 { + return true + } + + for _, breakdown := range usage.UsageBreakdownList { + if breakdown.CurrentUsageWithPrecision < breakdown.UsageLimitWithPrecision { + return false + } + + if breakdown.FreeTrialInfo != nil { + if breakdown.FreeTrialInfo.CurrentUsageWithPrecision < breakdown.FreeTrialInfo.UsageLimitWithPrecision { + return false + } + } + } + + return true +} + +// GetQuotaStatus retrieves a comprehensive quota status for a token. +func (c *UsageChecker) GetQuotaStatus(ctx context.Context, tokenData *KiroTokenData) (*QuotaStatus, error) { + usage, err := c.CheckUsage(ctx, tokenData) + if err != nil { + return nil, err + } + + status := &QuotaStatus{ + IsExhausted: IsQuotaExhausted(usage), + } + + if len(usage.UsageBreakdownList) > 0 { + breakdown := usage.UsageBreakdownList[0] + status.TotalLimit = breakdown.UsageLimitWithPrecision + status.CurrentUsage = breakdown.CurrentUsageWithPrecision + status.RemainingQuota = breakdown.UsageLimitWithPrecision - breakdown.CurrentUsageWithPrecision + status.ResourceType = breakdown.ResourceType + + if breakdown.FreeTrialInfo != nil { + status.TotalLimit += breakdown.FreeTrialInfo.UsageLimitWithPrecision + status.CurrentUsage += breakdown.FreeTrialInfo.CurrentUsageWithPrecision + freeRemaining := breakdown.FreeTrialInfo.UsageLimitWithPrecision - breakdown.FreeTrialInfo.CurrentUsageWithPrecision + if freeRemaining > 0 { + status.RemainingQuota += freeRemaining + } + } + } + + if usage.NextDateReset > 0 { + status.NextReset = time.Unix(int64(usage.NextDateReset/1000), 0) + } + + return status, nil +} + +// CalculateAvailableCount calculates the available request count based on usage limits. +func CalculateAvailableCount(usage *UsageQuotaResponse) float64 { + return GetRemainingQuota(usage) +} + +// GetUsagePercentage calculates the usage percentage. +func GetUsagePercentage(usage *UsageQuotaResponse) float64 { + if usage == nil || len(usage.UsageBreakdownList) == 0 { + return 100.0 + } + + var totalLimit, totalUsage float64 + for _, breakdown := range usage.UsageBreakdownList { + totalLimit += breakdown.UsageLimitWithPrecision + totalUsage += breakdown.CurrentUsageWithPrecision + + if breakdown.FreeTrialInfo != nil { + totalLimit += breakdown.FreeTrialInfo.UsageLimitWithPrecision + totalUsage += breakdown.FreeTrialInfo.CurrentUsageWithPrecision + } + } + + if totalLimit == 0 { + return 100.0 + } + + return (totalUsage / totalLimit) * 100 +} diff --git a/internal/registry/kiro_model_converter.go b/internal/registry/kiro_model_converter.go new file mode 100644 index 00000000..fe50a8f3 --- /dev/null +++ b/internal/registry/kiro_model_converter.go @@ -0,0 +1,303 @@ +// Package registry provides Kiro model conversion utilities. +// This file handles converting dynamic Kiro API model lists to the internal ModelInfo format, +// and merging with static metadata for thinking support and other capabilities. +package registry + +import ( + "strings" + "time" +) + +// KiroAPIModel represents a model from Kiro API response. +// This is a local copy to avoid import cycles with the kiro package. +// The structure mirrors kiro.KiroModel for easy data conversion. +type KiroAPIModel struct { + // ModelID is the unique identifier for the model (e.g., "claude-sonnet-4.5") + ModelID string + // ModelName is the human-readable name + ModelName string + // Description is the model description + Description string + // RateMultiplier is the credit multiplier for this model + RateMultiplier float64 + // RateUnit is the unit for rate calculation (e.g., "credit") + RateUnit string + // MaxInputTokens is the maximum input token limit + MaxInputTokens int +} + +// DefaultKiroThinkingSupport defines the default thinking configuration for Kiro models. +// All Kiro models support thinking with the following budget range. +var DefaultKiroThinkingSupport = &ThinkingSupport{ + Min: 1024, // Minimum thinking budget tokens + Max: 32000, // Maximum thinking budget tokens + ZeroAllowed: true, // Allow disabling thinking with 0 + DynamicAllowed: true, // Allow dynamic thinking budget (-1) +} + +// DefaultKiroContextLength is the default context window size for Kiro models. +const DefaultKiroContextLength = 200000 + +// DefaultKiroMaxCompletionTokens is the default max completion tokens for Kiro models. +const DefaultKiroMaxCompletionTokens = 64000 + +// ConvertKiroAPIModels converts Kiro API models to internal ModelInfo format. +// It performs the following transformations: +// - Normalizes model ID (e.g., claude-sonnet-4.5 → kiro-claude-sonnet-4-5) +// - Adds default thinking support metadata +// - Sets default context length and max completion tokens if not provided +// +// Parameters: +// - kiroModels: List of models from Kiro API response +// +// Returns: +// - []*ModelInfo: Converted model information list +func ConvertKiroAPIModels(kiroModels []*KiroAPIModel) []*ModelInfo { + if len(kiroModels) == 0 { + return nil + } + + now := time.Now().Unix() + result := make([]*ModelInfo, 0, len(kiroModels)) + + for _, km := range kiroModels { + // Skip nil models + if km == nil { + continue + } + + // Skip models without valid ID + if km.ModelID == "" { + continue + } + + // Normalize the model ID to kiro-* format + normalizedID := normalizeKiroModelID(km.ModelID) + + // Create ModelInfo with converted data + info := &ModelInfo{ + ID: normalizedID, + Object: "model", + Created: now, + OwnedBy: "aws", + Type: "kiro", + DisplayName: generateKiroDisplayName(km.ModelName, normalizedID), + Description: km.Description, + // Use MaxInputTokens from API if available, otherwise use default + ContextLength: getContextLength(km.MaxInputTokens), + MaxCompletionTokens: DefaultKiroMaxCompletionTokens, + // All Kiro models support thinking + Thinking: cloneThinkingSupport(DefaultKiroThinkingSupport), + } + + result = append(result, info) + } + + return result +} + +// GenerateAgenticVariants creates -agentic variants for each model. +// Agentic variants are optimized for coding agents with chunked writes. +// +// Parameters: +// - models: Base models to generate variants for +// +// Returns: +// - []*ModelInfo: Combined list of base models and their agentic variants +func GenerateAgenticVariants(models []*ModelInfo) []*ModelInfo { + if len(models) == 0 { + return nil + } + + // Pre-allocate result with capacity for both base models and variants + result := make([]*ModelInfo, 0, len(models)*2) + + for _, model := range models { + if model == nil { + continue + } + + // Add the base model first + result = append(result, model) + + // Skip if model already has -agentic suffix + if strings.HasSuffix(model.ID, "-agentic") { + continue + } + + // Skip special models that shouldn't have agentic variants + if model.ID == "kiro-auto" { + continue + } + + // Create agentic variant + agenticModel := &ModelInfo{ + ID: model.ID + "-agentic", + Object: model.Object, + Created: model.Created, + OwnedBy: model.OwnedBy, + Type: model.Type, + DisplayName: model.DisplayName + " (Agentic)", + Description: generateAgenticDescription(model.Description), + ContextLength: model.ContextLength, + MaxCompletionTokens: model.MaxCompletionTokens, + Thinking: cloneThinkingSupport(model.Thinking), + } + + result = append(result, agenticModel) + } + + return result +} + +// MergeWithStaticMetadata merges dynamic models with static metadata. +// Static metadata takes priority for any overlapping fields. +// This allows manual overrides for specific models while keeping dynamic discovery. +// +// Parameters: +// - dynamicModels: Models from Kiro API (converted to ModelInfo) +// - staticModels: Predefined model metadata (from GetKiroModels()) +// +// Returns: +// - []*ModelInfo: Merged model list with static metadata taking priority +func MergeWithStaticMetadata(dynamicModels, staticModels []*ModelInfo) []*ModelInfo { + if len(dynamicModels) == 0 && len(staticModels) == 0 { + return nil + } + + // Build a map of static models for quick lookup + staticMap := make(map[string]*ModelInfo, len(staticModels)) + for _, sm := range staticModels { + if sm != nil && sm.ID != "" { + staticMap[sm.ID] = sm + } + } + + // Build result, preferring static metadata where available + seenIDs := make(map[string]struct{}) + result := make([]*ModelInfo, 0, len(dynamicModels)+len(staticModels)) + + // First, process dynamic models and merge with static if available + for _, dm := range dynamicModels { + if dm == nil || dm.ID == "" { + continue + } + + // Skip duplicates + if _, seen := seenIDs[dm.ID]; seen { + continue + } + seenIDs[dm.ID] = struct{}{} + + // Check if static metadata exists for this model + if sm, exists := staticMap[dm.ID]; exists { + // Static metadata takes priority - use static model + result = append(result, sm) + } else { + // No static metadata - use dynamic model + result = append(result, dm) + } + } + + // Add any static models not in dynamic list + for _, sm := range staticModels { + if sm == nil || sm.ID == "" { + continue + } + if _, seen := seenIDs[sm.ID]; seen { + continue + } + seenIDs[sm.ID] = struct{}{} + result = append(result, sm) + } + + return result +} + +// normalizeKiroModelID converts Kiro API model IDs to internal format. +// Transformation rules: +// - Adds "kiro-" prefix if not present +// - Replaces dots with hyphens (e.g., 4.5 → 4-5) +// - Handles special cases like "auto" → "kiro-auto" +// +// Examples: +// - "claude-sonnet-4.5" → "kiro-claude-sonnet-4-5" +// - "claude-opus-4.5" → "kiro-claude-opus-4-5" +// - "auto" → "kiro-auto" +// - "kiro-claude-sonnet-4-5" → "kiro-claude-sonnet-4-5" (unchanged) +func normalizeKiroModelID(modelID string) string { + if modelID == "" { + return "" + } + + // Trim whitespace + modelID = strings.TrimSpace(modelID) + + // Replace dots with hyphens (e.g., 4.5 → 4-5) + normalized := strings.ReplaceAll(modelID, ".", "-") + + // Add kiro- prefix if not present + if !strings.HasPrefix(normalized, "kiro-") { + normalized = "kiro-" + normalized + } + + return normalized +} + +// generateKiroDisplayName creates a human-readable display name. +// Uses the API-provided model name if available, otherwise generates from ID. +func generateKiroDisplayName(modelName, normalizedID string) string { + if modelName != "" { + return "Kiro " + modelName + } + + // Generate from normalized ID by removing kiro- prefix and formatting + displayID := strings.TrimPrefix(normalizedID, "kiro-") + // Capitalize first letter of each word + words := strings.Split(displayID, "-") + for i, word := range words { + if len(word) > 0 { + words[i] = strings.ToUpper(word[:1]) + word[1:] + } + } + return "Kiro " + strings.Join(words, " ") +} + +// generateAgenticDescription creates description for agentic variants. +func generateAgenticDescription(baseDescription string) string { + if baseDescription == "" { + return "Optimized for coding agents with chunked writes" + } + return baseDescription + " (Agentic mode: chunked writes)" +} + +// getContextLength returns the context length, using default if not provided. +func getContextLength(maxInputTokens int) int { + if maxInputTokens > 0 { + return maxInputTokens + } + return DefaultKiroContextLength +} + +// cloneThinkingSupport creates a deep copy of ThinkingSupport. +// Returns nil if input is nil. +func cloneThinkingSupport(ts *ThinkingSupport) *ThinkingSupport { + if ts == nil { + return nil + } + + clone := &ThinkingSupport{ + Min: ts.Min, + Max: ts.Max, + ZeroAllowed: ts.ZeroAllowed, + DynamicAllowed: ts.DynamicAllowed, + } + + // Deep copy Levels slice if present + if len(ts.Levels) > 0 { + clone.Levels = make([]string, len(ts.Levels)) + copy(clone.Levels, ts.Levels) + } + + return clone +} diff --git a/internal/runtime/executor/kiro_executor.go b/internal/runtime/executor/kiro_executor.go index 3d152955..b0c14c61 100644 --- a/internal/runtime/executor/kiro_executor.go +++ b/internal/runtime/executor/kiro_executor.go @@ -7,13 +7,16 @@ import ( "encoding/base64" "encoding/binary" "encoding/json" + "errors" "fmt" "io" + "net" "net/http" "os" "path/filepath" "strings" "sync" + "syscall" "time" "github.com/google/uuid" @@ -53,9 +56,28 @@ const ( kiroIDEUserAgent = "aws-sdk-js/1.0.18 ua/2.1 os/darwin#25.0.0 lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.18 m/E KiroIDE-0.2.13-66c23a8c5d15afabec89ef9954ef52a119f10d369df04d548fc6c1eac694b0d1" kiroIDEAmzUserAgent = "aws-sdk-js/1.0.18 KiroIDE-0.2.13-66c23a8c5d15afabec89ef9954ef52a119f10d369df04d548fc6c1eac694b0d1" kiroIDEAgentModeSpec = "spec" - kiroAgentModeVibe = "vibe" + + // Socket retry configuration constants (based on kiro2Api reference implementation) + // Maximum number of retry attempts for socket/network errors + kiroSocketMaxRetries = 3 + // Base delay between retry attempts (uses exponential backoff: delay * 2^attempt) + kiroSocketBaseRetryDelay = 1 * time.Second + // Maximum delay between retry attempts (cap for exponential backoff) + kiroSocketMaxRetryDelay = 30 * time.Second + // First token timeout for streaming responses (how long to wait for first response) + kiroFirstTokenTimeout = 15 * time.Second + // Streaming read timeout (how long to wait between chunks) + kiroStreamingReadTimeout = 300 * time.Second ) +// retryableHTTPStatusCodes defines HTTP status codes that are considered retryable. +// Based on kiro2Api reference: 502 (Bad Gateway), 503 (Service Unavailable), 504 (Gateway Timeout) +var retryableHTTPStatusCodes = map[int]bool{ + 502: true, // Bad Gateway - upstream server error + 503: true, // Service Unavailable - server temporarily overloaded + 504: true, // Gateway Timeout - upstream server timeout +} + // Real-time usage estimation configuration // These control how often usage updates are sent during streaming var ( @@ -63,6 +85,241 @@ var ( usageUpdateTimeInterval = 15 * time.Second // Or every 15 seconds, whichever comes first ) +// Global FingerprintManager for dynamic User-Agent generation per token +// Each token gets a unique fingerprint on first use, which is cached for subsequent requests +var ( + globalFingerprintManager *kiroauth.FingerprintManager + globalFingerprintManagerOnce sync.Once +) + +// getGlobalFingerprintManager returns the global FingerprintManager instance +func getGlobalFingerprintManager() *kiroauth.FingerprintManager { + globalFingerprintManagerOnce.Do(func() { + globalFingerprintManager = kiroauth.NewFingerprintManager() + log.Infof("kiro: initialized global FingerprintManager for dynamic UA generation") + }) + return globalFingerprintManager +} + +// retryConfig holds configuration for socket retry logic. +// Based on kiro2Api Python implementation patterns. +type retryConfig struct { + MaxRetries int // Maximum number of retry attempts + BaseDelay time.Duration // Base delay between retries (exponential backoff) + MaxDelay time.Duration // Maximum delay cap + RetryableErrors []string // List of retryable error patterns + RetryableStatus map[int]bool // HTTP status codes to retry + FirstTokenTmout time.Duration // Timeout for first token in streaming + StreamReadTmout time.Duration // Timeout between stream chunks +} + +// defaultRetryConfig returns the default retry configuration for Kiro socket operations. +func defaultRetryConfig() retryConfig { + return retryConfig{ + MaxRetries: kiroSocketMaxRetries, + BaseDelay: kiroSocketBaseRetryDelay, + MaxDelay: kiroSocketMaxRetryDelay, + RetryableStatus: retryableHTTPStatusCodes, + RetryableErrors: []string{ + "connection reset", + "connection refused", + "broken pipe", + "EOF", + "timeout", + "temporary failure", + "no such host", + "network is unreachable", + "i/o timeout", + }, + FirstTokenTmout: kiroFirstTokenTimeout, + StreamReadTmout: kiroStreamingReadTimeout, + } +} + +// isRetryableError checks if an error is retryable based on error type and message. +// Returns true for network timeouts, connection resets, and temporary failures. +// Based on kiro2Api's retry logic patterns. +func isRetryableError(err error) bool { + if err == nil { + return false + } + + // Check for context cancellation - not retryable + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + + // Check for net.Error (timeout, temporary) + var netErr net.Error + if errors.As(err, &netErr) { + if netErr.Timeout() { + log.Debugf("kiro: isRetryableError: network timeout detected") + return true + } + // Note: Temporary() is deprecated but still useful for some error types + } + + // Check for specific syscall errors (connection reset, broken pipe, etc.) + var syscallErr syscall.Errno + if errors.As(err, &syscallErr) { + switch syscallErr { + case syscall.ECONNRESET: // Connection reset by peer + log.Debugf("kiro: isRetryableError: ECONNRESET detected") + return true + case syscall.ECONNREFUSED: // Connection refused + log.Debugf("kiro: isRetryableError: ECONNREFUSED detected") + return true + case syscall.EPIPE: // Broken pipe + log.Debugf("kiro: isRetryableError: EPIPE (broken pipe) detected") + return true + case syscall.ETIMEDOUT: // Connection timed out + log.Debugf("kiro: isRetryableError: ETIMEDOUT detected") + return true + case syscall.ENETUNREACH: // Network is unreachable + log.Debugf("kiro: isRetryableError: ENETUNREACH detected") + return true + case syscall.EHOSTUNREACH: // No route to host + log.Debugf("kiro: isRetryableError: EHOSTUNREACH detected") + return true + } + } + + // Check for net.OpError wrapping other errors + var opErr *net.OpError + if errors.As(err, &opErr) { + log.Debugf("kiro: isRetryableError: net.OpError detected, op=%s", opErr.Op) + // Recursively check the wrapped error + if opErr.Err != nil { + return isRetryableError(opErr.Err) + } + return true + } + + // Check error message for retryable patterns + errMsg := strings.ToLower(err.Error()) + cfg := defaultRetryConfig() + for _, pattern := range cfg.RetryableErrors { + if strings.Contains(errMsg, pattern) { + log.Debugf("kiro: isRetryableError: pattern '%s' matched in error: %s", pattern, errMsg) + return true + } + } + + // Check for EOF which may indicate connection was closed + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + log.Debugf("kiro: isRetryableError: EOF/UnexpectedEOF detected") + return true + } + + return false +} + +// isRetryableHTTPStatus checks if an HTTP status code is retryable. +// Based on kiro2Api: 502, 503, 504 are retryable server errors. +func isRetryableHTTPStatus(statusCode int) bool { + return retryableHTTPStatusCodes[statusCode] +} + +// calculateRetryDelay calculates the delay for the next retry attempt using exponential backoff. +// delay = min(baseDelay * 2^attempt, maxDelay) +// Adds ±30% jitter to prevent thundering herd. +func calculateRetryDelay(attempt int, cfg retryConfig) time.Duration { + return kiroauth.ExponentialBackoffWithJitter(attempt, cfg.BaseDelay, cfg.MaxDelay) +} + +// logRetryAttempt logs a retry attempt with relevant context. +func logRetryAttempt(attempt, maxRetries int, reason string, delay time.Duration, endpoint string) { + log.Warnf("kiro: retry attempt %d/%d for %s, waiting %v before next attempt (endpoint: %s)", + attempt+1, maxRetries, reason, delay, endpoint) +} + +// kiroHTTPClientPool provides a shared HTTP client with connection pooling for Kiro API. +// This reduces connection overhead and improves performance for concurrent requests. +// Based on kiro2Api's connection pooling pattern. +var ( + kiroHTTPClientPool *http.Client + kiroHTTPClientPoolOnce sync.Once +) + +// getKiroPooledHTTPClient returns a shared HTTP client with optimized connection pooling. +// The client is lazily initialized on first use and reused across requests. +// This is especially beneficial for: +// - Reducing TCP handshake overhead +// - Enabling HTTP/2 multiplexing +// - Better handling of keep-alive connections +func getKiroPooledHTTPClient() *http.Client { + kiroHTTPClientPoolOnce.Do(func() { + transport := &http.Transport{ + // Connection pool settings + MaxIdleConns: 100, // Max idle connections across all hosts + MaxIdleConnsPerHost: 20, // Max idle connections per host + MaxConnsPerHost: 50, // Max total connections per host + IdleConnTimeout: 90 * time.Second, // How long idle connections stay in pool + + // Timeouts for connection establishment + DialContext: (&net.Dialer{ + Timeout: 30 * time.Second, // TCP connection timeout + KeepAlive: 30 * time.Second, // TCP keep-alive interval + }).DialContext, + + // TLS handshake timeout + TLSHandshakeTimeout: 10 * time.Second, + + // Response header timeout + ResponseHeaderTimeout: 30 * time.Second, + + // Expect 100-continue timeout + ExpectContinueTimeout: 1 * time.Second, + + // Enable HTTP/2 when available + ForceAttemptHTTP2: true, + } + + kiroHTTPClientPool = &http.Client{ + Transport: transport, + // No global timeout - let individual requests set their own timeouts via context + } + + log.Debugf("kiro: initialized pooled HTTP client (MaxIdleConns=%d, MaxIdleConnsPerHost=%d, MaxConnsPerHost=%d)", + transport.MaxIdleConns, transport.MaxIdleConnsPerHost, transport.MaxConnsPerHost) + }) + + return kiroHTTPClientPool +} + +// newKiroHTTPClientWithPooling creates an HTTP client that uses connection pooling when appropriate. +// It respects proxy configuration from auth or config, falling back to the pooled client. +// This provides the best of both worlds: custom proxy support + connection reuse. +func newKiroHTTPClientWithPooling(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client { + // Check if a proxy is configured - if so, we need a custom client + var proxyURL string + if auth != nil { + proxyURL = strings.TrimSpace(auth.ProxyURL) + } + if proxyURL == "" && cfg != nil { + proxyURL = strings.TrimSpace(cfg.ProxyURL) + } + + // If proxy is configured, use the existing proxy-aware client (doesn't pool) + if proxyURL != "" { + log.Debugf("kiro: using proxy-aware HTTP client (proxy=%s)", proxyURL) + return newProxyAwareHTTPClient(ctx, cfg, auth, timeout) + } + + // No proxy - use pooled client for better performance + pooledClient := getKiroPooledHTTPClient() + + // If timeout is specified, we need to wrap the pooled transport with timeout + if timeout > 0 { + return &http.Client{ + Transport: pooledClient.Transport, + Timeout: timeout, + } + } + + return pooledClient +} + // kiroEndpointConfig bundles endpoint URL with its compatible Origin and AmzTarget values. // This solves the "triple mismatch" problem where different endpoints require matching // Origin and X-Amz-Target header values. @@ -99,7 +356,7 @@ var kiroEndpointConfigs = []kiroEndpointConfig{ Name: "CodeWhisperer", }, { - URL: "https://q.us-east-1.amazonaws.com/generateAssistantResponse", + URL: "https://q.us-east-1.amazonaws.com/", Origin: "CLI", AmzTarget: "AmazonQDeveloperStreamingService.SendMessage", Name: "AmazonQ", @@ -217,6 +474,29 @@ func NewKiroExecutor(cfg *config.Config) *KiroExecutor { // Identifier returns the unique identifier for this executor. func (e *KiroExecutor) Identifier() string { return "kiro" } +// applyDynamicFingerprint applies token-specific fingerprint headers to the request +// For IDC auth, uses dynamic fingerprint-based User-Agent +// For other auth types, uses static Amazon Q CLI style headers +func applyDynamicFingerprint(req *http.Request, auth *cliproxyauth.Auth) { + if isIDCAuth(auth) { + // Get token-specific fingerprint for dynamic UA generation + tokenKey := getTokenKey(auth) + fp := getGlobalFingerprintManager().GetFingerprint(tokenKey) + + // Use fingerprint-generated dynamic User-Agent + req.Header.Set("User-Agent", fp.BuildUserAgent()) + req.Header.Set("X-Amz-User-Agent", fp.BuildAmzUserAgent()) + req.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec) + + log.Debugf("kiro: using dynamic fingerprint for token %s (SDK:%s, OS:%s/%s, Kiro:%s)", + tokenKey[:8]+"...", fp.SDKVersion, fp.OSType, fp.OSVersion, fp.KiroVersion) + } else { + // Use static Amazon Q CLI style headers for non-IDC auth + req.Header.Set("User-Agent", kiroUserAgent) + req.Header.Set("X-Amz-User-Agent", kiroFullUserAgent) + } +} + // PrepareRequest prepares the HTTP request before execution. func (e *KiroExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { if req == nil { @@ -226,16 +506,10 @@ func (e *KiroExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth if strings.TrimSpace(accessToken) == "" { return statusErr{code: http.StatusUnauthorized, msg: "missing access token"} } - if isIDCAuth(auth) { - req.Header.Set("User-Agent", kiroIDEUserAgent) - req.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent) - req.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec) - } else { - req.Header.Set("User-Agent", kiroUserAgent) - req.Header.Set("X-Amz-User-Agent", kiroFullUserAgent) - req.Header.Set("x-amzn-kiro-agent-mode", kiroAgentModeVibe) - } - req.Header.Set("x-amzn-codewhisperer-optout", "true") + + // Apply dynamic fingerprint-based headers + applyDynamicFingerprint(req, auth) + req.Header.Set("Amz-Sdk-Request", "attempt=1; max=3") req.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String()) req.Header.Set("Authorization", "Bearer "+accessToken) @@ -259,10 +533,23 @@ func (e *KiroExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, if errPrepare := e.PrepareRequest(httpReq, auth); errPrepare != nil { return nil, errPrepare } - httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + httpClient := newKiroHTTPClientWithPooling(ctx, e.cfg, auth, 0) return httpClient.Do(httpReq) } +// getTokenKey returns a unique key for rate limiting based on auth credentials. +// Uses auth ID if available, otherwise falls back to a hash of the access token. +func getTokenKey(auth *cliproxyauth.Auth) string { + if auth != nil && auth.ID != "" { + return auth.ID + } + accessToken, _ := kiroCredentials(auth) + if len(accessToken) > 16 { + return accessToken[:16] + } + return accessToken +} + // Execute sends the request to Kiro API and returns the response. // Supports automatic token refresh on 401/403 errors. func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { @@ -271,6 +558,24 @@ func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req return resp, fmt.Errorf("kiro: access token not found in auth") } + // Rate limiting: get token key for tracking + tokenKey := getTokenKey(auth) + rateLimiter := kiroauth.GetGlobalRateLimiter() + cooldownMgr := kiroauth.GetGlobalCooldownManager() + + // Check if token is in cooldown period + if cooldownMgr.IsInCooldown(tokenKey) { + remaining := cooldownMgr.GetRemainingCooldown(tokenKey) + reason := cooldownMgr.GetCooldownReason(tokenKey) + log.Warnf("kiro: token %s is in cooldown (reason: %s), remaining: %v", tokenKey, reason, remaining) + return resp, fmt.Errorf("kiro: token is in cooldown for %v (reason: %s)", remaining, reason) + } + + // Wait for rate limiter before proceeding + log.Debugf("kiro: waiting for rate limiter for token %s", tokenKey) + rateLimiter.WaitForToken(tokenKey) + log.Debugf("kiro: rate limiter cleared for token %s", tokenKey) + reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth) defer reporter.trackFailure(ctx, &err) @@ -303,7 +608,7 @@ func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req // Execute with retry on 401/403 and 429 (quota exhausted) // Note: currentOrigin and kiroPayload are built inside executeWithRetry for each endpoint - resp, err = e.executeWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, nil, body, from, to, reporter, "", kiroModelID, isAgentic, isChatOnly) + resp, err = e.executeWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, nil, body, from, to, reporter, "", kiroModelID, isAgentic, isChatOnly, tokenKey) return resp, err } @@ -312,9 +617,12 @@ func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req // - Amazon Q endpoint (CLI origin) uses Amazon Q Developer quota // - CodeWhisperer endpoint (AI_EDITOR origin) uses Kiro IDE quota // Also supports multi-endpoint fallback similar to Antigravity implementation. -func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, accessToken, profileArn string, kiroPayload, body []byte, from, to sdktranslator.Format, reporter *usageReporter, currentOrigin, kiroModelID string, isAgentic, isChatOnly bool) (cliproxyexecutor.Response, error) { +// tokenKey is used for rate limiting and cooldown tracking. +func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, accessToken, profileArn string, kiroPayload, body []byte, from, to sdktranslator.Format, reporter *usageReporter, currentOrigin, kiroModelID string, isAgentic, isChatOnly bool, tokenKey string) (cliproxyexecutor.Response, error) { var resp cliproxyexecutor.Response maxRetries := 2 // Allow retries for token refresh + endpoint fallback + rateLimiter := kiroauth.GetGlobalRateLimiter() + cooldownMgr := kiroauth.GetGlobalCooldownManager() endpointConfigs := getKiroEndpointConfigs(auth) var last429Err error @@ -332,6 +640,12 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. endpointIdx+1, len(endpointConfigs), url, endpointConfig.Name, currentOrigin) for attempt := 0; attempt <= maxRetries; attempt++ { + // Apply human-like delay before first request (not on retries) + // This mimics natural user behavior patterns + if attempt == 0 && endpointIdx == 0 { + kiroauth.ApplyHumanLikeDelay() + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroPayload)) if err != nil { return resp, err @@ -342,20 +656,9 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. // Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors) httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget) - // Use different headers based on auth type - // IDC auth uses Kiro IDE style headers (from kiro2api) - // Other auth types use Amazon Q CLI style headers - if isIDCAuth(auth) { - httpReq.Header.Set("User-Agent", kiroIDEUserAgent) - httpReq.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent) - httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec) - log.Debugf("kiro: using Kiro IDE headers for IDC auth") - } else { - httpReq.Header.Set("User-Agent", kiroUserAgent) - httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent) - httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroAgentModeVibe) - } - httpReq.Header.Set("x-amzn-codewhisperer-optout", "true") + // Apply dynamic fingerprint-based headers + applyDynamicFingerprint(httpReq, auth) + httpReq.Header.Set("Amz-Sdk-Request", "attempt=1; max=3") httpReq.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String()) @@ -386,10 +689,34 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. AuthValue: authValue, }) - httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 120*time.Second) + httpClient := newKiroHTTPClientWithPooling(ctx, e.cfg, auth, 120*time.Second) httpResp, err := httpClient.Do(httpReq) if err != nil { + // Check for context cancellation first - client disconnected, not a server error + // Use 499 (Client Closed Request - nginx convention) instead of 500 + if errors.Is(err, context.Canceled) { + log.Debugf("kiro: request canceled by client (context.Canceled)") + return resp, statusErr{code: 499, msg: "client canceled request"} + } + + // Check for context deadline exceeded - request timed out + // Return 504 Gateway Timeout instead of 500 + if errors.Is(err, context.DeadlineExceeded) { + log.Debugf("kiro: request timed out (context.DeadlineExceeded)") + return resp, statusErr{code: http.StatusGatewayTimeout, msg: "upstream request timed out"} + } + recordAPIResponseError(ctx, e.cfg, err) + + // Enhanced socket retry: Check if error is retryable (network timeout, connection reset, etc.) + retryCfg := defaultRetryConfig() + if isRetryableError(err) && attempt < retryCfg.MaxRetries { + delay := calculateRetryDelay(attempt, retryCfg) + logRetryAttempt(attempt, retryCfg.MaxRetries, fmt.Sprintf("socket error: %v", err), delay, endpointConfig.Name) + time.Sleep(delay) + continue + } + return resp, err } recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) @@ -401,6 +728,12 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. _ = httpResp.Body.Close() appendAPIResponseChunk(ctx, e.cfg, respBody) + // Record failure and set cooldown for 429 + rateLimiter.MarkTokenFailed(tokenKey) + cooldownDuration := kiroauth.CalculateCooldownFor429(attempt) + cooldownMgr.SetCooldown(tokenKey, cooldownDuration, kiroauth.CooldownReason429) + log.Warnf("kiro: rate limit hit (429), token %s set to cooldown for %v", tokenKey, cooldownDuration) + // Preserve last 429 so callers can correctly backoff when all endpoints are exhausted last429Err = statusErr{code: httpResp.StatusCode, msg: string(respBody)} @@ -412,13 +745,21 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. } // Handle 5xx server errors with exponential backoff retry + // Enhanced: Use retryConfig for consistent retry behavior if httpResp.StatusCode >= 500 && httpResp.StatusCode < 600 { respBody, _ := io.ReadAll(httpResp.Body) _ = httpResp.Body.Close() appendAPIResponseChunk(ctx, e.cfg, respBody) - if attempt < maxRetries { - // Exponential backoff: 1s, 2s, 4s... (max 30s) + retryCfg := defaultRetryConfig() + // Check if this specific 5xx code is retryable (502, 503, 504) + if isRetryableHTTPStatus(httpResp.StatusCode) && attempt < retryCfg.MaxRetries { + delay := calculateRetryDelay(attempt, retryCfg) + logRetryAttempt(attempt, retryCfg.MaxRetries, fmt.Sprintf("HTTP %d", httpResp.StatusCode), delay, endpointConfig.Name) + time.Sleep(delay) + continue + } else if attempt < maxRetries { + // Fallback for other 5xx errors (500, 501, etc.) backoff := time.Duration(1< 30*time.Second { backoff = 30 * time.Second @@ -492,7 +833,10 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. // Check for SUSPENDED status - return immediately without retry if strings.Contains(respBodyStr, "SUSPENDED") || strings.Contains(respBodyStr, "TEMPORARILY_SUSPENDED") { - log.Errorf("kiro: account is suspended, cannot proceed") + // Set long cooldown for suspended accounts + rateLimiter.CheckAndMarkSuspended(tokenKey, respBodyStr) + cooldownMgr.SetCooldown(tokenKey, kiroauth.LongCooldown, kiroauth.CooldownReasonSuspended) + log.Errorf("kiro: account is suspended, token %s set to cooldown for %v", tokenKey, kiroauth.LongCooldown) return resp, statusErr{code: httpResp.StatusCode, msg: "account suspended: " + string(respBody)} } @@ -581,6 +925,10 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. appendAPIResponseChunk(ctx, e.cfg, []byte(content)) reporter.publish(ctx, usageInfo) + // Record success for rate limiting + rateLimiter.MarkTokenSuccess(tokenKey) + log.Debugf("kiro: request successful, token %s marked as success", tokenKey) + // Build response in Claude format for Kiro translator // stopReason is extracted from upstream response by parseEventStream kiroResponse := kiroclaude.BuildClaudeResponse(content, toolUses, req.Model, usageInfo, stopReason) @@ -608,6 +956,24 @@ func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut return nil, fmt.Errorf("kiro: access token not found in auth") } + // Rate limiting: get token key for tracking + tokenKey := getTokenKey(auth) + rateLimiter := kiroauth.GetGlobalRateLimiter() + cooldownMgr := kiroauth.GetGlobalCooldownManager() + + // Check if token is in cooldown period + if cooldownMgr.IsInCooldown(tokenKey) { + remaining := cooldownMgr.GetRemainingCooldown(tokenKey) + reason := cooldownMgr.GetCooldownReason(tokenKey) + log.Warnf("kiro: token %s is in cooldown (reason: %s), remaining: %v", tokenKey, reason, remaining) + return nil, fmt.Errorf("kiro: token is in cooldown for %v (reason: %s)", remaining, reason) + } + + // Wait for rate limiter before proceeding + log.Debugf("kiro: stream waiting for rate limiter for token %s", tokenKey) + rateLimiter.WaitForToken(tokenKey) + log.Debugf("kiro: stream rate limiter cleared for token %s", tokenKey) + reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth) defer reporter.trackFailure(ctx, &err) @@ -640,7 +1006,7 @@ func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut // Execute stream with retry on 401/403 and 429 (quota exhausted) // Note: currentOrigin and kiroPayload are built inside executeStreamWithRetry for each endpoint - return e.executeStreamWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, nil, body, from, reporter, "", kiroModelID, isAgentic, isChatOnly) + return e.executeStreamWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, nil, body, from, reporter, "", kiroModelID, isAgentic, isChatOnly, tokenKey) } // executeStreamWithRetry performs the streaming HTTP request with automatic retry on auth errors. @@ -648,8 +1014,11 @@ func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut // - Amazon Q endpoint (CLI origin) uses Amazon Q Developer quota // - CodeWhisperer endpoint (AI_EDITOR origin) uses Kiro IDE quota // Also supports multi-endpoint fallback similar to Antigravity implementation. -func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, accessToken, profileArn string, kiroPayload, body []byte, from sdktranslator.Format, reporter *usageReporter, currentOrigin, kiroModelID string, isAgentic, isChatOnly bool) (<-chan cliproxyexecutor.StreamChunk, error) { +// tokenKey is used for rate limiting and cooldown tracking. +func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, accessToken, profileArn string, kiroPayload, body []byte, from sdktranslator.Format, reporter *usageReporter, currentOrigin, kiroModelID string, isAgentic, isChatOnly bool, tokenKey string) (<-chan cliproxyexecutor.StreamChunk, error) { maxRetries := 2 // Allow retries for token refresh + endpoint fallback + rateLimiter := kiroauth.GetGlobalRateLimiter() + cooldownMgr := kiroauth.GetGlobalCooldownManager() endpointConfigs := getKiroEndpointConfigs(auth) var last429Err error @@ -667,6 +1036,13 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox endpointIdx+1, len(endpointConfigs), url, endpointConfig.Name, currentOrigin) for attempt := 0; attempt <= maxRetries; attempt++ { + // Apply human-like delay before first streaming request (not on retries) + // This mimics natural user behavior patterns + // Note: Delay is NOT applied during streaming response - only before initial request + if attempt == 0 && endpointIdx == 0 { + kiroauth.ApplyHumanLikeDelay() + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroPayload)) if err != nil { return nil, err @@ -677,20 +1053,9 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox // Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors) httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget) - // Use different headers based on auth type - // IDC auth uses Kiro IDE style headers (from kiro2api) - // Other auth types use Amazon Q CLI style headers - if isIDCAuth(auth) { - httpReq.Header.Set("User-Agent", kiroIDEUserAgent) - httpReq.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent) - httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec) - log.Debugf("kiro: using Kiro IDE headers for IDC auth") - } else { - httpReq.Header.Set("User-Agent", kiroUserAgent) - httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent) - httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroAgentModeVibe) - } - httpReq.Header.Set("x-amzn-codewhisperer-optout", "true") + // Apply dynamic fingerprint-based headers + applyDynamicFingerprint(httpReq, auth) + httpReq.Header.Set("Amz-Sdk-Request", "attempt=1; max=3") httpReq.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String()) @@ -721,10 +1086,20 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox AuthValue: authValue, }) - httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + httpClient := newKiroHTTPClientWithPooling(ctx, e.cfg, auth, 0) httpResp, err := httpClient.Do(httpReq) if err != nil { recordAPIResponseError(ctx, e.cfg, err) + + // Enhanced socket retry for streaming: Check if error is retryable (network timeout, connection reset, etc.) + retryCfg := defaultRetryConfig() + if isRetryableError(err) && attempt < retryCfg.MaxRetries { + delay := calculateRetryDelay(attempt, retryCfg) + logRetryAttempt(attempt, retryCfg.MaxRetries, fmt.Sprintf("stream socket error: %v", err), delay, endpointConfig.Name) + time.Sleep(delay) + continue + } + return nil, err } recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) @@ -736,6 +1111,12 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox _ = httpResp.Body.Close() appendAPIResponseChunk(ctx, e.cfg, respBody) + // Record failure and set cooldown for 429 + rateLimiter.MarkTokenFailed(tokenKey) + cooldownDuration := kiroauth.CalculateCooldownFor429(attempt) + cooldownMgr.SetCooldown(tokenKey, cooldownDuration, kiroauth.CooldownReason429) + log.Warnf("kiro: stream rate limit hit (429), token %s set to cooldown for %v", tokenKey, cooldownDuration) + // Preserve last 429 so callers can correctly backoff when all endpoints are exhausted last429Err = statusErr{code: httpResp.StatusCode, msg: string(respBody)} @@ -747,13 +1128,21 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox } // Handle 5xx server errors with exponential backoff retry + // Enhanced: Use retryConfig for consistent retry behavior if httpResp.StatusCode >= 500 && httpResp.StatusCode < 600 { respBody, _ := io.ReadAll(httpResp.Body) _ = httpResp.Body.Close() appendAPIResponseChunk(ctx, e.cfg, respBody) - if attempt < maxRetries { - // Exponential backoff: 1s, 2s, 4s... (max 30s) + retryCfg := defaultRetryConfig() + // Check if this specific 5xx code is retryable (502, 503, 504) + if isRetryableHTTPStatus(httpResp.StatusCode) && attempt < retryCfg.MaxRetries { + delay := calculateRetryDelay(attempt, retryCfg) + logRetryAttempt(attempt, retryCfg.MaxRetries, fmt.Sprintf("stream HTTP %d", httpResp.StatusCode), delay, endpointConfig.Name) + time.Sleep(delay) + continue + } else if attempt < maxRetries { + // Fallback for other 5xx errors (500, 501, etc.) backoff := time.Duration(1< 30*time.Second { backoff = 30 * time.Second @@ -840,7 +1229,10 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox // Check for SUSPENDED status - return immediately without retry if strings.Contains(respBodyStr, "SUSPENDED") || strings.Contains(respBodyStr, "TEMPORARILY_SUSPENDED") { - log.Errorf("kiro: account is suspended, cannot proceed") + // Set long cooldown for suspended accounts + rateLimiter.CheckAndMarkSuspended(tokenKey, respBodyStr) + cooldownMgr.SetCooldown(tokenKey, kiroauth.LongCooldown, kiroauth.CooldownReasonSuspended) + log.Errorf("kiro: stream account is suspended, token %s set to cooldown for %v", tokenKey, kiroauth.LongCooldown) return nil, statusErr{code: httpResp.StatusCode, msg: "account suspended: " + string(respBody)} } @@ -890,6 +1282,11 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox out := make(chan cliproxyexecutor.StreamChunk) + // Record success immediately since connection was established successfully + // Streaming errors will be handled separately + rateLimiter.MarkTokenSuccess(tokenKey) + log.Debugf("kiro: stream request successful, token %s marked as success", tokenKey) + go func(resp *http.Response, thinkingEnabled bool) { defer close(out) defer func() { diff --git a/internal/translator/kiro/common/utf8_stream.go b/internal/translator/kiro/common/utf8_stream.go new file mode 100644 index 00000000..b8d24c82 --- /dev/null +++ b/internal/translator/kiro/common/utf8_stream.go @@ -0,0 +1,97 @@ +package common + +import ( + "unicode/utf8" +) + +type UTF8StreamParser struct { + buffer []byte +} + +func NewUTF8StreamParser() *UTF8StreamParser { + return &UTF8StreamParser{ + buffer: make([]byte, 0, 64), + } +} + +func (p *UTF8StreamParser) Write(data []byte) { + p.buffer = append(p.buffer, data...) +} + +func (p *UTF8StreamParser) Read() (string, bool) { + if len(p.buffer) == 0 { + return "", false + } + + validLen := p.findValidUTF8End(p.buffer) + if validLen == 0 { + return "", false + } + + result := string(p.buffer[:validLen]) + p.buffer = p.buffer[validLen:] + + return result, true +} + +func (p *UTF8StreamParser) Flush() string { + if len(p.buffer) == 0 { + return "" + } + result := string(p.buffer) + p.buffer = p.buffer[:0] + return result +} + +func (p *UTF8StreamParser) Reset() { + p.buffer = p.buffer[:0] +} + +func (p *UTF8StreamParser) findValidUTF8End(data []byte) int { + if len(data) == 0 { + return 0 + } + + end := len(data) + for i := 1; i <= 3 && i <= len(data); i++ { + b := data[len(data)-i] + if b&0x80 == 0 { + break + } + if b&0xC0 == 0xC0 { + size := p.utf8CharSize(b) + available := i + if size > available { + end = len(data) - i + } + break + } + } + + if end > 0 && !utf8.Valid(data[:end]) { + for i := end - 1; i >= 0; i-- { + if utf8.Valid(data[:i+1]) { + return i + 1 + } + } + return 0 + } + + return end +} + +func (p *UTF8StreamParser) utf8CharSize(b byte) int { + if b&0x80 == 0 { + return 1 + } + if b&0xE0 == 0xC0 { + return 2 + } + if b&0xF0 == 0xE0 { + return 3 + } + if b&0xF8 == 0xF0 { + return 4 + } + return 1 +} diff --git a/internal/translator/kiro/common/utf8_stream_test.go b/internal/translator/kiro/common/utf8_stream_test.go new file mode 100644 index 00000000..23e80989 --- /dev/null +++ b/internal/translator/kiro/common/utf8_stream_test.go @@ -0,0 +1,402 @@ +package common + +import ( + "strings" + "sync" + "testing" + "unicode/utf8" +) + +func TestNewUTF8StreamParser(t *testing.T) { + p := NewUTF8StreamParser() + if p == nil { + t.Fatal("expected non-nil UTF8StreamParser") + } + if p.buffer == nil { + t.Error("expected non-nil buffer") + } +} + +func TestWrite(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("hello")) + + result, ok := p.Read() + if !ok { + t.Error("expected ok to be true") + } + if result != "hello" { + t.Errorf("expected 'hello', got '%s'", result) + } +} + +func TestWrite_MultipleWrites(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("hel")) + p.Write([]byte("lo")) + + result, ok := p.Read() + if !ok { + t.Error("expected ok to be true") + } + if result != "hello" { + t.Errorf("expected 'hello', got '%s'", result) + } +} + +func TestRead_EmptyBuffer(t *testing.T) { + p := NewUTF8StreamParser() + result, ok := p.Read() + if ok { + t.Error("expected ok to be false for empty buffer") + } + if result != "" { + t.Errorf("expected empty string, got '%s'", result) + } +} + +func TestRead_IncompleteUTF8(t *testing.T) { + p := NewUTF8StreamParser() + + // Write incomplete multi-byte UTF-8 character + // 中 (U+4E2D) = E4 B8 AD + p.Write([]byte{0xE4, 0xB8}) + + result, ok := p.Read() + if ok { + t.Error("expected ok to be false for incomplete UTF-8") + } + if result != "" { + t.Errorf("expected empty string, got '%s'", result) + } + + // Complete the character + p.Write([]byte{0xAD}) + result, ok = p.Read() + if !ok { + t.Error("expected ok to be true after completing UTF-8") + } + if result != "中" { + t.Errorf("expected '中', got '%s'", result) + } +} + +func TestRead_MixedASCIIAndUTF8(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("Hello 世界")) + + result, ok := p.Read() + if !ok { + t.Error("expected ok to be true") + } + if result != "Hello 世界" { + t.Errorf("expected 'Hello 世界', got '%s'", result) + } +} + +func TestRead_PartialMultibyteAtEnd(t *testing.T) { + p := NewUTF8StreamParser() + // "Hello" + partial "世" (E4 B8 96) + p.Write([]byte("Hello")) + p.Write([]byte{0xE4, 0xB8}) + + result, ok := p.Read() + if !ok { + t.Error("expected ok to be true for valid portion") + } + if result != "Hello" { + t.Errorf("expected 'Hello', got '%s'", result) + } + + // Complete the character + p.Write([]byte{0x96}) + result, ok = p.Read() + if !ok { + t.Error("expected ok to be true after completing") + } + if result != "世" { + t.Errorf("expected '世', got '%s'", result) + } +} + +func TestFlush(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("hello")) + + result := p.Flush() + if result != "hello" { + t.Errorf("expected 'hello', got '%s'", result) + } + + // Verify buffer is cleared + result2, ok := p.Read() + if ok { + t.Error("expected ok to be false after flush") + } + if result2 != "" { + t.Errorf("expected empty string after flush, got '%s'", result2) + } +} + +func TestFlush_EmptyBuffer(t *testing.T) { + p := NewUTF8StreamParser() + result := p.Flush() + if result != "" { + t.Errorf("expected empty string, got '%s'", result) + } +} + +func TestFlush_IncompleteUTF8(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte{0xE4, 0xB8}) + + result := p.Flush() + // Flush returns everything including incomplete bytes + if len(result) != 2 { + t.Errorf("expected 2 bytes flushed, got %d", len(result)) + } +} + +func TestReset(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("hello")) + p.Reset() + + result, ok := p.Read() + if ok { + t.Error("expected ok to be false after reset") + } + if result != "" { + t.Errorf("expected empty string after reset, got '%s'", result) + } +} + +func TestUtf8CharSize(t *testing.T) { + p := NewUTF8StreamParser() + + testCases := []struct { + b byte + expected int + }{ + {0x00, 1}, // ASCII + {0x7F, 1}, // ASCII max + {0xC0, 2}, // 2-byte start + {0xDF, 2}, // 2-byte start + {0xE0, 3}, // 3-byte start + {0xEF, 3}, // 3-byte start + {0xF0, 4}, // 4-byte start + {0xF7, 4}, // 4-byte start + {0x80, 1}, // Continuation byte (fallback) + } + + for _, tc := range testCases { + size := p.utf8CharSize(tc.b) + if size != tc.expected { + t.Errorf("utf8CharSize(0x%X) = %d, expected %d", tc.b, size, tc.expected) + } + } +} + +func TestStreamingScenario(t *testing.T) { + p := NewUTF8StreamParser() + + // Simulate streaming: "Hello, 世界! 🌍" + chunks := [][]byte{ + []byte("Hello, "), + {0xE4, 0xB8}, // partial 世 + {0x96, 0xE7}, // complete 世, partial 界 + {0x95, 0x8C}, // complete 界 + []byte("! "), + {0xF0, 0x9F}, // partial 🌍 + {0x8C, 0x8D}, // complete 🌍 + } + + var results []string + for _, chunk := range chunks { + p.Write(chunk) + if result, ok := p.Read(); ok { + results = append(results, result) + } + } + + combined := strings.Join(results, "") + if combined != "Hello, 世界! 🌍" { + t.Errorf("expected 'Hello, 世界! 🌍', got '%s'", combined) + } +} + +func TestValidUTF8Output(t *testing.T) { + p := NewUTF8StreamParser() + + testStrings := []string{ + "Hello World", + "你好世界", + "こんにちは", + "🎉🎊🎁", + "Mixed 混合 Текст ტექსტი", + } + + for _, s := range testStrings { + p.Reset() + p.Write([]byte(s)) + result, ok := p.Read() + if !ok { + t.Errorf("expected ok for '%s'", s) + } + if !utf8.ValidString(result) { + t.Errorf("invalid UTF-8 output for input '%s'", s) + } + if result != s { + t.Errorf("expected '%s', got '%s'", s, result) + } + } +} + +func TestLargeData(t *testing.T) { + p := NewUTF8StreamParser() + + // Generate large UTF-8 string + var builder strings.Builder + for i := 0; i < 1000; i++ { + builder.WriteString("Hello 世界! ") + } + largeString := builder.String() + + p.Write([]byte(largeString)) + result, ok := p.Read() + if !ok { + t.Error("expected ok for large data") + } + if result != largeString { + t.Error("large data mismatch") + } +} + +func TestByteByByteWriting(t *testing.T) { + p := NewUTF8StreamParser() + input := "Hello 世界" + inputBytes := []byte(input) + + var results []string + for _, b := range inputBytes { + p.Write([]byte{b}) + if result, ok := p.Read(); ok { + results = append(results, result) + } + } + + combined := strings.Join(results, "") + if combined != input { + t.Errorf("expected '%s', got '%s'", input, combined) + } +} + +func TestEmoji4ByteUTF8(t *testing.T) { + p := NewUTF8StreamParser() + + // 🎉 = F0 9F 8E 89 + emoji := "🎉" + emojiBytes := []byte(emoji) + + for i := 0; i < len(emojiBytes)-1; i++ { + p.Write(emojiBytes[i : i+1]) + result, ok := p.Read() + if ok && result != "" { + t.Errorf("unexpected output before emoji complete: '%s'", result) + } + } + + p.Write(emojiBytes[len(emojiBytes)-1:]) + result, ok := p.Read() + if !ok { + t.Error("expected ok after completing emoji") + } + if result != emoji { + t.Errorf("expected '%s', got '%s'", emoji, result) + } +} + +func TestContinuationBytesOnly(t *testing.T) { + p := NewUTF8StreamParser() + + // Write only continuation bytes (invalid UTF-8) + p.Write([]byte{0x80, 0x80, 0x80}) + + result, ok := p.Read() + // Should handle gracefully - either return nothing or return the bytes + _ = result + _ = ok +} + +func TestUTF8StreamParser_ConcurrentSafety(t *testing.T) { + // Note: UTF8StreamParser doesn't have built-in locks, + // so this test verifies it works with external synchronization + p := NewUTF8StreamParser() + var mu sync.Mutex + const numGoroutines = 10 + const numOperations = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func() { + defer wg.Done() + for j := 0; j < numOperations; j++ { + mu.Lock() + switch j % 4 { + case 0: + p.Write([]byte("test")) + case 1: + p.Read() + case 2: + p.Flush() + case 3: + p.Reset() + } + mu.Unlock() + } + }() + } + + wg.Wait() +} + +func TestConsecutiveReads(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("hello")) + + result1, ok1 := p.Read() + if !ok1 || result1 != "hello" { + t.Error("first read failed") + } + + result2, ok2 := p.Read() + if ok2 || result2 != "" { + t.Error("second read should return empty") + } +} + +func TestFlushThenWrite(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("first")) + p.Flush() + p.Write([]byte("second")) + + result, ok := p.Read() + if !ok || result != "second" { + t.Errorf("expected 'second', got '%s'", result) + } +} + +func TestResetThenWrite(t *testing.T) { + p := NewUTF8StreamParser() + p.Write([]byte("first")) + p.Reset() + p.Write([]byte("second")) + + result, ok := p.Read() + if !ok || result != "second" { + t.Errorf("expected 'second', got '%s'", result) + } +} diff --git a/internal/watcher/events.go b/internal/watcher/events.go index eb428353..fb96ad2a 100644 --- a/internal/watcher/events.go +++ b/internal/watcher/events.go @@ -170,7 +170,9 @@ func (w *Watcher) handleKiroIDETokenChange(event fsnotify.Event) { } } - tokenData, err := kiroauth.LoadKiroIDEToken() + // Use retry logic to handle file lock contention (e.g., Kiro IDE writing the file) + // This prevents "being used by another process" errors on Windows + tokenData, err := kiroauth.LoadKiroIDETokenWithRetry(10, 50*time.Millisecond) if err != nil { log.Debugf("failed to load Kiro IDE token after change: %v", err) return diff --git a/sdk/auth/kiro.go b/sdk/auth/kiro.go index b75cd28e..7747c777 100644 --- a/sdk/auth/kiro.go +++ b/sdk/auth/kiro.go @@ -12,9 +12,9 @@ import ( ) // extractKiroIdentifier extracts a meaningful identifier for file naming. -// Returns account name if provided, otherwise profile ARN ID. +// Returns account name if provided, otherwise profile ARN ID, then client ID. // All extracted values are sanitized to prevent path injection attacks. -func extractKiroIdentifier(accountName, profileArn string) string { +func extractKiroIdentifier(accountName, profileArn, clientID string) string { // Priority 1: Use account name if provided if accountName != "" { return kiroauth.SanitizeEmailForFilename(accountName) @@ -29,6 +29,11 @@ func extractKiroIdentifier(accountName, profileArn string) string { } } + // Priority 3: Use client ID (for IDC auth without email/profileArn) + if clientID != "" { + return kiroauth.SanitizeEmailForFilename(clientID) + } + // Fallback: timestamp return fmt.Sprintf("%d", time.Now().UnixNano()%100000) } @@ -62,7 +67,7 @@ func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, } // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn, tokenData.ClientID) // Determine label based on auth method label := fmt.Sprintf("kiro-%s", source) @@ -173,7 +178,7 @@ func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.C } // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn, tokenData.ClientID) now := time.Now() fileName := fmt.Sprintf("kiro-aws-%s.json", idPart) @@ -217,129 +222,17 @@ func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.C } // LoginWithGoogle performs OAuth login for Kiro with Google. -// This uses a custom protocol handler (kiro://) to receive the callback. +// NOTE: Google login is not available for third-party applications due to AWS Cognito restrictions. +// Please use AWS Builder ID or import your token from Kiro IDE. func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { - if cfg == nil { - return nil, fmt.Errorf("kiro auth: configuration is required") - } - - oauth := kiroauth.NewKiroOAuth(cfg) - - // Use Google OAuth flow with protocol handler - tokenData, err := oauth.LoginWithGoogle(ctx) - if err != nil { - return nil, fmt.Errorf("google login failed: %w", err) - } - - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) - - now := time.Now() - fileName := fmt.Sprintf("kiro-google-%s.json", idPart) - - record := &coreauth.Auth{ - ID: fileName, - Provider: "kiro", - FileName: fileName, - Label: "kiro-google", - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Metadata: map[string]any{ - "type": "kiro", - "access_token": tokenData.AccessToken, - "refresh_token": tokenData.RefreshToken, - "profile_arn": tokenData.ProfileArn, - "expires_at": tokenData.ExpiresAt, - "auth_method": tokenData.AuthMethod, - "provider": tokenData.Provider, - "email": tokenData.Email, - }, - Attributes: map[string]string{ - "profile_arn": tokenData.ProfileArn, - "source": "google-oauth", - "email": tokenData.Email, - }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), - } - - if tokenData.Email != "" { - fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email) - } else { - fmt.Println("\n✓ Kiro Google authentication completed successfully!") - } - - return record, nil + return nil, fmt.Errorf("Google login is not available for third-party applications due to AWS Cognito restrictions.\n\nAlternatives:\n 1. Use AWS Builder ID: cliproxy kiro --builder-id\n 2. Import token from Kiro IDE: cliproxy kiro --import\n\nTo get a token from Kiro IDE:\n 1. Open Kiro IDE and login with Google\n 2. Find: ~/.kiro/kiro-auth-token.json\n 3. Run: cliproxy kiro --import") } // LoginWithGitHub performs OAuth login for Kiro with GitHub. -// This uses a custom protocol handler (kiro://) to receive the callback. +// NOTE: GitHub login is not available for third-party applications due to AWS Cognito restrictions. +// Please use AWS Builder ID or import your token from Kiro IDE. func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { - if cfg == nil { - return nil, fmt.Errorf("kiro auth: configuration is required") - } - - oauth := kiroauth.NewKiroOAuth(cfg) - - // Use GitHub OAuth flow with protocol handler - tokenData, err := oauth.LoginWithGitHub(ctx) - if err != nil { - return nil, fmt.Errorf("github login failed: %w", err) - } - - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) - - now := time.Now() - fileName := fmt.Sprintf("kiro-github-%s.json", idPart) - - record := &coreauth.Auth{ - ID: fileName, - Provider: "kiro", - FileName: fileName, - Label: "kiro-github", - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Metadata: map[string]any{ - "type": "kiro", - "access_token": tokenData.AccessToken, - "refresh_token": tokenData.RefreshToken, - "profile_arn": tokenData.ProfileArn, - "expires_at": tokenData.ExpiresAt, - "auth_method": tokenData.AuthMethod, - "provider": tokenData.Provider, - "email": tokenData.Email, - }, - Attributes: map[string]string{ - "profile_arn": tokenData.ProfileArn, - "source": "github-oauth", - "email": tokenData.Email, - }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), - } - - if tokenData.Email != "" { - fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email) - } else { - fmt.Println("\n✓ Kiro GitHub authentication completed successfully!") - } - - return record, nil + return nil, fmt.Errorf("GitHub login is not available for third-party applications due to AWS Cognito restrictions.\n\nAlternatives:\n 1. Use AWS Builder ID: cliproxy kiro --builder-id\n 2. Import token from Kiro IDE: cliproxy kiro --import\n\nTo get a token from Kiro IDE:\n 1. Open Kiro IDE and login with GitHub\n 2. Find: ~/.kiro/kiro-auth-token.json\n 3. Run: cliproxy kiro --import") } // ImportFromKiroIDE imports token from Kiro IDE's token file. @@ -361,7 +254,7 @@ func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.C } // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn, tokenData.ClientID) // Sanitize provider to prevent path traversal (defense-in-depth) provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider))) if provider == "" { @@ -387,12 +280,17 @@ func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.C "expires_at": tokenData.ExpiresAt, "auth_method": tokenData.AuthMethod, "provider": tokenData.Provider, + "client_id": tokenData.ClientID, + "client_secret": tokenData.ClientSecret, "email": tokenData.Email, + "region": tokenData.Region, + "start_url": tokenData.StartURL, }, Attributes: map[string]string{ "profile_arn": tokenData.ProfileArn, "source": "kiro-ide-import", "email": tokenData.Email, + "region": tokenData.Region, }, // NextRefreshAfter is aligned with RefreshLead (5min) NextRefreshAfter: expiresAt.Add(-5 * time.Minute), diff --git a/sdk/auth/kiro.go.bak b/sdk/auth/kiro.go.bak new file mode 100644 index 00000000..b75cd28e --- /dev/null +++ b/sdk/auth/kiro.go.bak @@ -0,0 +1,470 @@ +package auth + +import ( + "context" + "fmt" + "strings" + "time" + + kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro" + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" +) + +// extractKiroIdentifier extracts a meaningful identifier for file naming. +// Returns account name if provided, otherwise profile ARN ID. +// All extracted values are sanitized to prevent path injection attacks. +func extractKiroIdentifier(accountName, profileArn string) string { + // Priority 1: Use account name if provided + if accountName != "" { + return kiroauth.SanitizeEmailForFilename(accountName) + } + + // Priority 2: Use profile ARN ID part (sanitized to prevent path injection) + if profileArn != "" { + parts := strings.Split(profileArn, "/") + if len(parts) >= 2 { + // Sanitize the ARN component to prevent path traversal + return kiroauth.SanitizeEmailForFilename(parts[len(parts)-1]) + } + } + + // Fallback: timestamp + return fmt.Sprintf("%d", time.Now().UnixNano()%100000) +} + +// KiroAuthenticator implements OAuth authentication for Kiro with Google login. +type KiroAuthenticator struct{} + +// NewKiroAuthenticator constructs a Kiro authenticator. +func NewKiroAuthenticator() *KiroAuthenticator { + return &KiroAuthenticator{} +} + +// Provider returns the provider key for the authenticator. +func (a *KiroAuthenticator) Provider() string { + return "kiro" +} + +// RefreshLead indicates how soon before expiry a refresh should be attempted. +// Set to 5 minutes to match Antigravity and avoid frequent refresh checks while still ensuring timely token refresh. +func (a *KiroAuthenticator) RefreshLead() *time.Duration { + d := 5 * time.Minute + return &d +} + +// createAuthRecord creates an auth record from token data. +func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, source string) (*coreauth.Auth, error) { + // Parse expires_at + expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) + if err != nil { + expiresAt = time.Now().Add(1 * time.Hour) + } + + // Extract identifier for file naming + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + + // Determine label based on auth method + label := fmt.Sprintf("kiro-%s", source) + if tokenData.AuthMethod == "idc" { + label = "kiro-idc" + } + + now := time.Now() + fileName := fmt.Sprintf("%s-%s.json", label, idPart) + + metadata := map[string]any{ + "type": "kiro", + "access_token": tokenData.AccessToken, + "refresh_token": tokenData.RefreshToken, + "profile_arn": tokenData.ProfileArn, + "expires_at": tokenData.ExpiresAt, + "auth_method": tokenData.AuthMethod, + "provider": tokenData.Provider, + "client_id": tokenData.ClientID, + "client_secret": tokenData.ClientSecret, + "email": tokenData.Email, + } + + // Add IDC-specific fields if present + if tokenData.StartURL != "" { + metadata["start_url"] = tokenData.StartURL + } + if tokenData.Region != "" { + metadata["region"] = tokenData.Region + } + + attributes := map[string]string{ + "profile_arn": tokenData.ProfileArn, + "source": source, + "email": tokenData.Email, + } + + // Add IDC-specific attributes if present + if tokenData.AuthMethod == "idc" { + attributes["source"] = "aws-idc" + if tokenData.StartURL != "" { + attributes["start_url"] = tokenData.StartURL + } + if tokenData.Region != "" { + attributes["region"] = tokenData.Region + } + } + + record := &coreauth.Auth{ + ID: fileName, + Provider: "kiro", + FileName: fileName, + Label: label, + Status: coreauth.StatusActive, + CreatedAt: now, + UpdatedAt: now, + Metadata: metadata, + Attributes: attributes, + // NextRefreshAfter is aligned with RefreshLead (5min) + NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + } + + if tokenData.Email != "" { + fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email) + } else { + fmt.Println("\n✓ Kiro authentication completed successfully!") + } + + return record, nil +} + +// Login performs OAuth login for Kiro with AWS (Builder ID or IDC). +// This shows a method selection prompt and handles both flows. +func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { + if cfg == nil { + return nil, fmt.Errorf("kiro auth: configuration is required") + } + + // Use the unified method selection flow (Builder ID or IDC) + ssoClient := kiroauth.NewSSOOIDCClient(cfg) + tokenData, err := ssoClient.LoginWithMethodSelection(ctx) + if err != nil { + return nil, fmt.Errorf("login failed: %w", err) + } + + return a.createAuthRecord(tokenData, "aws") +} + +// LoginWithAuthCode performs OAuth login for Kiro with AWS Builder ID using authorization code flow. +// This provides a better UX than device code flow as it uses automatic browser callback. +func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { + if cfg == nil { + return nil, fmt.Errorf("kiro auth: configuration is required") + } + + oauth := kiroauth.NewKiroOAuth(cfg) + + // Use AWS Builder ID authorization code flow + tokenData, err := oauth.LoginWithBuilderIDAuthCode(ctx) + if err != nil { + return nil, fmt.Errorf("login failed: %w", err) + } + + // Parse expires_at + expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) + if err != nil { + expiresAt = time.Now().Add(1 * time.Hour) + } + + // Extract identifier for file naming + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + + now := time.Now() + fileName := fmt.Sprintf("kiro-aws-%s.json", idPart) + + record := &coreauth.Auth{ + ID: fileName, + Provider: "kiro", + FileName: fileName, + Label: "kiro-aws", + Status: coreauth.StatusActive, + CreatedAt: now, + UpdatedAt: now, + Metadata: map[string]any{ + "type": "kiro", + "access_token": tokenData.AccessToken, + "refresh_token": tokenData.RefreshToken, + "profile_arn": tokenData.ProfileArn, + "expires_at": tokenData.ExpiresAt, + "auth_method": tokenData.AuthMethod, + "provider": tokenData.Provider, + "client_id": tokenData.ClientID, + "client_secret": tokenData.ClientSecret, + "email": tokenData.Email, + }, + Attributes: map[string]string{ + "profile_arn": tokenData.ProfileArn, + "source": "aws-builder-id-authcode", + "email": tokenData.Email, + }, + // NextRefreshAfter is aligned with RefreshLead (5min) + NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + } + + if tokenData.Email != "" { + fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email) + } else { + fmt.Println("\n✓ Kiro authentication completed successfully!") + } + + return record, nil +} + +// LoginWithGoogle performs OAuth login for Kiro with Google. +// This uses a custom protocol handler (kiro://) to receive the callback. +func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { + if cfg == nil { + return nil, fmt.Errorf("kiro auth: configuration is required") + } + + oauth := kiroauth.NewKiroOAuth(cfg) + + // Use Google OAuth flow with protocol handler + tokenData, err := oauth.LoginWithGoogle(ctx) + if err != nil { + return nil, fmt.Errorf("google login failed: %w", err) + } + + // Parse expires_at + expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) + if err != nil { + expiresAt = time.Now().Add(1 * time.Hour) + } + + // Extract identifier for file naming + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + + now := time.Now() + fileName := fmt.Sprintf("kiro-google-%s.json", idPart) + + record := &coreauth.Auth{ + ID: fileName, + Provider: "kiro", + FileName: fileName, + Label: "kiro-google", + Status: coreauth.StatusActive, + CreatedAt: now, + UpdatedAt: now, + Metadata: map[string]any{ + "type": "kiro", + "access_token": tokenData.AccessToken, + "refresh_token": tokenData.RefreshToken, + "profile_arn": tokenData.ProfileArn, + "expires_at": tokenData.ExpiresAt, + "auth_method": tokenData.AuthMethod, + "provider": tokenData.Provider, + "email": tokenData.Email, + }, + Attributes: map[string]string{ + "profile_arn": tokenData.ProfileArn, + "source": "google-oauth", + "email": tokenData.Email, + }, + // NextRefreshAfter is aligned with RefreshLead (5min) + NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + } + + if tokenData.Email != "" { + fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email) + } else { + fmt.Println("\n✓ Kiro Google authentication completed successfully!") + } + + return record, nil +} + +// LoginWithGitHub performs OAuth login for Kiro with GitHub. +// This uses a custom protocol handler (kiro://) to receive the callback. +func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { + if cfg == nil { + return nil, fmt.Errorf("kiro auth: configuration is required") + } + + oauth := kiroauth.NewKiroOAuth(cfg) + + // Use GitHub OAuth flow with protocol handler + tokenData, err := oauth.LoginWithGitHub(ctx) + if err != nil { + return nil, fmt.Errorf("github login failed: %w", err) + } + + // Parse expires_at + expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) + if err != nil { + expiresAt = time.Now().Add(1 * time.Hour) + } + + // Extract identifier for file naming + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + + now := time.Now() + fileName := fmt.Sprintf("kiro-github-%s.json", idPart) + + record := &coreauth.Auth{ + ID: fileName, + Provider: "kiro", + FileName: fileName, + Label: "kiro-github", + Status: coreauth.StatusActive, + CreatedAt: now, + UpdatedAt: now, + Metadata: map[string]any{ + "type": "kiro", + "access_token": tokenData.AccessToken, + "refresh_token": tokenData.RefreshToken, + "profile_arn": tokenData.ProfileArn, + "expires_at": tokenData.ExpiresAt, + "auth_method": tokenData.AuthMethod, + "provider": tokenData.Provider, + "email": tokenData.Email, + }, + Attributes: map[string]string{ + "profile_arn": tokenData.ProfileArn, + "source": "github-oauth", + "email": tokenData.Email, + }, + // NextRefreshAfter is aligned with RefreshLead (5min) + NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + } + + if tokenData.Email != "" { + fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email) + } else { + fmt.Println("\n✓ Kiro GitHub authentication completed successfully!") + } + + return record, nil +} + +// ImportFromKiroIDE imports token from Kiro IDE's token file. +func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.Config) (*coreauth.Auth, error) { + tokenData, err := kiroauth.LoadKiroIDEToken() + if err != nil { + return nil, fmt.Errorf("failed to load Kiro IDE token: %w", err) + } + + // Parse expires_at + expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) + if err != nil { + expiresAt = time.Now().Add(1 * time.Hour) + } + + // Extract email from JWT if not already set (for imported tokens) + if tokenData.Email == "" { + tokenData.Email = kiroauth.ExtractEmailFromJWT(tokenData.AccessToken) + } + + // Extract identifier for file naming + idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) + // Sanitize provider to prevent path traversal (defense-in-depth) + provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider))) + if provider == "" { + provider = "imported" // Fallback for legacy tokens without provider + } + + now := time.Now() + fileName := fmt.Sprintf("kiro-%s-%s.json", provider, idPart) + + record := &coreauth.Auth{ + ID: fileName, + Provider: "kiro", + FileName: fileName, + Label: fmt.Sprintf("kiro-%s", provider), + Status: coreauth.StatusActive, + CreatedAt: now, + UpdatedAt: now, + Metadata: map[string]any{ + "type": "kiro", + "access_token": tokenData.AccessToken, + "refresh_token": tokenData.RefreshToken, + "profile_arn": tokenData.ProfileArn, + "expires_at": tokenData.ExpiresAt, + "auth_method": tokenData.AuthMethod, + "provider": tokenData.Provider, + "email": tokenData.Email, + }, + Attributes: map[string]string{ + "profile_arn": tokenData.ProfileArn, + "source": "kiro-ide-import", + "email": tokenData.Email, + }, + // NextRefreshAfter is aligned with RefreshLead (5min) + NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + } + + // Display the email if extracted + if tokenData.Email != "" { + fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s, Account: %s)\n", tokenData.Provider, tokenData.Email) + } else { + fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s)\n", tokenData.Provider) + } + + return record, nil +} + +// Refresh refreshes an expired Kiro token using AWS SSO OIDC. +func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, auth *coreauth.Auth) (*coreauth.Auth, error) { + if auth == nil || auth.Metadata == nil { + return nil, fmt.Errorf("invalid auth record") + } + + refreshToken, ok := auth.Metadata["refresh_token"].(string) + if !ok || refreshToken == "" { + return nil, fmt.Errorf("refresh token not found") + } + + clientID, _ := auth.Metadata["client_id"].(string) + clientSecret, _ := auth.Metadata["client_secret"].(string) + authMethod, _ := auth.Metadata["auth_method"].(string) + startURL, _ := auth.Metadata["start_url"].(string) + region, _ := auth.Metadata["region"].(string) + + var tokenData *kiroauth.KiroTokenData + var err error + + ssoClient := kiroauth.NewSSOOIDCClient(cfg) + + // Use SSO OIDC refresh for AWS Builder ID or IDC, otherwise use Kiro's OAuth refresh endpoint + switch { + case clientID != "" && clientSecret != "" && authMethod == "idc" && region != "": + // IDC refresh with region-specific endpoint + tokenData, err = ssoClient.RefreshTokenWithRegion(ctx, clientID, clientSecret, refreshToken, region, startURL) + case clientID != "" && clientSecret != "" && authMethod == "builder-id": + // Builder ID refresh with default endpoint + tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken) + default: + // Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub) + oauth := kiroauth.NewKiroOAuth(cfg) + tokenData, err = oauth.RefreshToken(ctx, refreshToken) + } + + if err != nil { + return nil, fmt.Errorf("token refresh failed: %w", err) + } + + // Parse expires_at + expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) + if err != nil { + expiresAt = time.Now().Add(1 * time.Hour) + } + + // Clone auth to avoid mutating the input parameter + updated := auth.Clone() + now := time.Now() + updated.UpdatedAt = now + updated.LastRefreshedAt = now + updated.Metadata["access_token"] = tokenData.AccessToken + updated.Metadata["refresh_token"] = tokenData.RefreshToken + updated.Metadata["expires_at"] = tokenData.ExpiresAt + updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization + // NextRefreshAfter is aligned with RefreshLead (5min) + updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute) + + return updated, nil +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 66d1b8dd..885304ad 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -13,6 +13,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v6/internal/api" + kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro" "github.com/router-for-me/CLIProxyAPI/v6/internal/registry" "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor" _ "github.com/router-for-me/CLIProxyAPI/v6/internal/usage" @@ -775,7 +776,7 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) { models = registry.GetGitHubCopilotModels() models = applyExcludedModels(models, excluded) case "kiro": - models = registry.GetKiroModels() + models = s.fetchKiroModels(a) models = applyExcludedModels(models, excluded) default: // Handle OpenAI-compatibility providers by name using config @@ -1338,3 +1339,201 @@ func applyOAuthModelAlias(cfg *config.Config, provider, authKind string, models } return out } + +// fetchKiroModels attempts to dynamically fetch Kiro models from the API. +// If dynamic fetch fails, it falls back to static registry.GetKiroModels(). +func (s *Service) fetchKiroModels(a *coreauth.Auth) []*ModelInfo { + if a == nil { + log.Debug("kiro: auth is nil, using static models") + return registry.GetKiroModels() + } + + // Extract token data from auth attributes + tokenData := s.extractKiroTokenData(a) + if tokenData == nil || tokenData.AccessToken == "" { + log.Debug("kiro: no valid token data in auth, using static models") + return registry.GetKiroModels() + } + + // Create KiroAuth instance + kAuth := kiroauth.NewKiroAuth(s.cfg) + if kAuth == nil { + log.Warn("kiro: failed to create KiroAuth instance, using static models") + return registry.GetKiroModels() + } + + // Use timeout context for API call + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + // Attempt to fetch dynamic models + apiModels, err := kAuth.ListAvailableModels(ctx, tokenData) + if err != nil { + log.Warnf("kiro: failed to fetch dynamic models: %v, using static models", err) + return registry.GetKiroModels() + } + + if len(apiModels) == 0 { + log.Debug("kiro: API returned no models, using static models") + return registry.GetKiroModels() + } + + // Convert API models to ModelInfo + models := convertKiroAPIModels(apiModels) + + // Generate agentic variants + models = generateKiroAgenticVariants(models) + + log.Infof("kiro: successfully fetched %d models from API (including agentic variants)", len(models)) + return models +} + +// extractKiroTokenData extracts KiroTokenData from auth attributes and metadata. +func (s *Service) extractKiroTokenData(a *coreauth.Auth) *kiroauth.KiroTokenData { + if a == nil || a.Attributes == nil { + return nil + } + + accessToken := strings.TrimSpace(a.Attributes["access_token"]) + if accessToken == "" { + return nil + } + + tokenData := &kiroauth.KiroTokenData{ + AccessToken: accessToken, + ProfileArn: strings.TrimSpace(a.Attributes["profile_arn"]), + } + + // Also try to get refresh token from metadata + if a.Metadata != nil { + if rt, ok := a.Metadata["refresh_token"].(string); ok { + tokenData.RefreshToken = rt + } + } + + return tokenData +} + +// convertKiroAPIModels converts Kiro API models to ModelInfo slice. +func convertKiroAPIModels(apiModels []*kiroauth.KiroModel) []*ModelInfo { + if len(apiModels) == 0 { + return nil + } + + now := time.Now().Unix() + models := make([]*ModelInfo, 0, len(apiModels)) + + for _, m := range apiModels { + if m == nil || m.ModelID == "" { + continue + } + + // Create model ID with kiro- prefix + modelID := "kiro-" + normalizeKiroModelID(m.ModelID) + + info := &ModelInfo{ + ID: modelID, + Object: "model", + Created: now, + OwnedBy: "aws", + Type: "kiro", + DisplayName: formatKiroDisplayName(m.ModelName, m.RateMultiplier), + Description: m.Description, + ContextLength: 200000, + MaxCompletionTokens: 64000, + Thinking: ®istry.ThinkingSupport{Min: 1024, Max: 32000, ZeroAllowed: true, DynamicAllowed: true}, + } + + if m.MaxInputTokens > 0 { + info.ContextLength = m.MaxInputTokens + } + + models = append(models, info) + } + + return models +} + +// normalizeKiroModelID normalizes a Kiro model ID by converting dots to dashes +// and removing common prefixes. +func normalizeKiroModelID(modelID string) string { + // Remove common prefixes + modelID = strings.TrimPrefix(modelID, "anthropic.") + modelID = strings.TrimPrefix(modelID, "amazon.") + + // Replace dots with dashes for consistency + modelID = strings.ReplaceAll(modelID, ".", "-") + + // Replace underscores with dashes + modelID = strings.ReplaceAll(modelID, "_", "-") + + return strings.ToLower(modelID) +} + +// formatKiroDisplayName formats the display name with rate multiplier info. +func formatKiroDisplayName(modelName string, rateMultiplier float64) string { + if modelName == "" { + return "" + } + + displayName := "Kiro " + modelName + if rateMultiplier > 0 && rateMultiplier != 1.0 { + displayName += fmt.Sprintf(" (%.1fx credit)", rateMultiplier) + } + + return displayName +} + +// generateKiroAgenticVariants generates agentic variants for Kiro models. +// Agentic variants have optimized system prompts for coding agents. +func generateKiroAgenticVariants(models []*ModelInfo) []*ModelInfo { + if len(models) == 0 { + return models + } + + result := make([]*ModelInfo, 0, len(models)*2) + result = append(result, models...) + + for _, m := range models { + if m == nil { + continue + } + + // Skip if already an agentic variant + if strings.HasSuffix(m.ID, "-agentic") { + continue + } + + // Skip auto models from agentic variant generation + if strings.Contains(m.ID, "-auto") { + continue + } + + // Create agentic variant + agentic := &ModelInfo{ + ID: m.ID + "-agentic", + Object: m.Object, + Created: m.Created, + OwnedBy: m.OwnedBy, + Type: m.Type, + DisplayName: m.DisplayName + " (Agentic)", + Description: m.Description + " - Optimized for coding agents (chunked writes)", + ContextLength: m.ContextLength, + MaxCompletionTokens: m.MaxCompletionTokens, + } + + // Copy thinking support if present + if m.Thinking != nil { + agentic.Thinking = ®istry.ThinkingSupport{ + Min: m.Thinking.Min, + Max: m.Thinking.Max, + ZeroAllowed: m.Thinking.ZeroAllowed, + DynamicAllowed: m.Thinking.DynamicAllowed, + } + } + + result = append(result, agentic) + } + + return result +} diff --git a/test_api.py b/test_api.py new file mode 100644 index 00000000..1849e2ba --- /dev/null +++ b/test_api.py @@ -0,0 +1,452 @@ +#!/usr/bin/env python3 +""" +CLIProxyAPI 全面测试脚本 +测试模型列表、流式输出、thinking模式及复杂任务 +""" + +import requests +import json +import time +import sys +import io +from typing import Optional, List, Dict, Any + +# 修复 Windows 控制台编码问题 +sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace') +sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace') + +# 配置 +BASE_URL = "http://localhost:8317" +API_KEY = "your-api-key-1" +HEADERS = { + "Authorization": f"Bearer {API_KEY}", + "Content-Type": "application/json" +} + +# 复杂任务提示词 - 用于测试 thinking 模式 +COMPLEX_TASK_PROMPT = """请帮我分析以下复杂的编程问题，并给出详细的解决方案： + +问题：设计一个高并发的分布式任务调度系统，需要满足以下要求： +1. 支持百万级任务队列 +2. 任务可以设置优先级、延迟执行、定时执行 +3. 支持任务依赖关系（DAG调度） +4. 失败重试机制，支持指数退避 +5. 任务结果持久化和查询 +6. 水平扩展能力 +7. 监控和告警 + +请从以下几个方面详细分析： +1. 整体架构设计 +2. 核心数据结构 +3. 调度算法选择 +4. 容错机制设计 +5. 性能优化策略 +6. 技术选型建议 + +请逐步思考每个方面，给出你的推理过程。""" + +# 简单测试提示词 +SIMPLE_PROMPT = "Hello! Please respond with 'OK' if you receive this message." + +def print_separator(title: str): + print(f"\n{'='*60}") + print(f" {title}") + print(f"{'='*60}\n") + +def print_result(name: str, success: bool, detail: str = ""): + status = "✅ PASS" if success else "❌ FAIL" + print(f"{status} | {name}") + if detail: + print(f" └─ {detail[:200]}{'...' if len(detail) > 200 else ''}") + +def get_models() -> List[str]: + """获取可用模型列表""" + print_separator("获取模型列表") + try: + resp = requests.get(f"{BASE_URL}/v1/models", headers=HEADERS, timeout=30) + if resp.status_code == 200: + data = resp.json() + models = [m.get("id", m.get("name", "unknown")) for m in data.get("data", [])] + print(f"找到 {len(models)} 个模型:") + for m in models: + print(f" - {m}") + return models + else: + print(f"❌ 获取模型列表失败: HTTP {resp.status_code}") + print(f" 响应: {resp.text[:500]}") + return [] + except Exception as e: + print(f"❌ 获取模型列表异常: {e}") + return [] + +def test_model_basic(model: str) -> tuple: + """基础可用性测试，返回 (success, error_detail)""" + try: + payload = { + "model": model, + "messages": [{"role": "user", "content": SIMPLE_PROMPT}], + "max_tokens": 50, + "stream": False + } + resp = requests.post( + f"{BASE_URL}/v1/chat/completions", + headers=HEADERS, + json=payload, + timeout=60 + ) + if resp.status_code == 200: + data = resp.json() + content = data.get("choices", [{}])[0].get("message", {}).get("content", "") + return (bool(content), f"content_len={len(content)}") + else: + return (False, f"HTTP {resp.status_code}: {resp.text[:300]}") + except Exception as e: + return (False, str(e)) + +def test_streaming(model: str) -> Dict[str, Any]: + """测试流式输出""" + result = {"success": False, "chunks": 0, "content": "", "error": None} + try: + payload = { + "model": model, + "messages": [{"role": "user", "content": "Count from 1 to 5, one number per line."}], + "max_tokens": 100, + "stream": True + } + resp = requests.post( + f"{BASE_URL}/v1/chat/completions", + headers=HEADERS, + json=payload, + timeout=60, + stream=True + ) + + if resp.status_code != 200: + result["error"] = f"HTTP {resp.status_code}: {resp.text[:200]}" + return result + + content_parts = [] + for line in resp.iter_lines(): + if line: + line_str = line.decode('utf-8') + if line_str.startswith("data: "): + data_str = line_str[6:] + if data_str.strip() == "[DONE]": + break + try: + data = json.loads(data_str) + result["chunks"] += 1 + choices = data.get("choices", []) + if choices: + delta = choices[0].get("delta", {}) + if "content" in delta and delta["content"]: + content_parts.append(delta["content"]) + except json.JSONDecodeError: + pass + except Exception as e: + result["error"] = f"Parse error: {e}, data: {data_str[:200]}" + + result["content"] = "".join(content_parts) + result["success"] = result["chunks"] > 0 and len(result["content"]) > 0 + + except Exception as e: + result["error"] = str(e) + + return result + +def test_thinking_mode(model: str, complex_task: bool = False) -> Dict[str, Any]: + """测试 thinking 模式""" + result = { + "success": False, + "has_reasoning": False, + "reasoning_content": "", + "content": "", + "error": None, + "chunks": 0 + } + + prompt = COMPLEX_TASK_PROMPT if complex_task else "What is 15 * 23? Please think step by step." + + try: + # 尝试不同的 thinking 模式参数格式 + payload = { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 8000 if complex_task else 2000, + "stream": True + } + + # 根据模型类型添加 thinking 参数 + if "claude" in model.lower(): + payload["thinking"] = {"type": "enabled", "budget_tokens": 5000 if complex_task else 2000} + elif "gemini" in model.lower(): + payload["thinking"] = {"thinking_budget": 5000 if complex_task else 2000} + elif "gpt" in model.lower() or "codex" in model.lower() or "o1" in model.lower() or "o3" in model.lower(): + payload["reasoning_effort"] = "high" if complex_task else "medium" + else: + # 通用格式 + payload["thinking"] = {"type": "enabled", "budget_tokens": 5000 if complex_task else 2000} + + resp = requests.post( + f"{BASE_URL}/v1/chat/completions", + headers=HEADERS, + json=payload, + timeout=300 if complex_task else 120, + stream=True + ) + + if resp.status_code != 200: + result["error"] = f"HTTP {resp.status_code}: {resp.text[:500]}" + return result + + content_parts = [] + reasoning_parts = [] + + for line in resp.iter_lines(): + if line: + line_str = line.decode('utf-8') + if line_str.startswith("data: "): + data_str = line_str[6:] + if data_str.strip() == "[DONE]": + break + try: + data = json.loads(data_str) + result["chunks"] += 1 + + choices = data.get("choices", []) + if not choices: + continue + choice = choices[0] + delta = choice.get("delta", {}) + + # 检查 reasoning_content (Claude/OpenAI格式) + if "reasoning_content" in delta and delta["reasoning_content"]: + reasoning_parts.append(delta["reasoning_content"]) + result["has_reasoning"] = True + + # 检查 thinking (Gemini格式) + if "thinking" in delta and delta["thinking"]: + reasoning_parts.append(delta["thinking"]) + result["has_reasoning"] = True + + # 常规内容 + if "content" in delta and delta["content"]: + content_parts.append(delta["content"]) + + except json.JSONDecodeError as e: + pass + except Exception as e: + result["error"] = f"Parse error: {e}" + + result["reasoning_content"] = "".join(reasoning_parts) + result["content"] = "".join(content_parts) + result["success"] = result["chunks"] > 0 and (len(result["content"]) > 0 or len(result["reasoning_content"]) > 0) + + except requests.exceptions.Timeout: + result["error"] = "Request timeout" + except Exception as e: + result["error"] = str(e) + + return result + +def run_full_test(): + """运行完整测试""" + print("\n" + "="*60) + print(" CLIProxyAPI 全面测试") + print("="*60) + print(f"目标地址: {BASE_URL}") + print(f"API Key: {API_KEY[:10]}...") + + # 1. 获取模型列表 + models = get_models() + if not models: + print("\n❌ 无法获取模型列表，测试终止") + return + + # 2. 基础可用性测试 + print_separator("基础可用性测试") + available_models = [] + for model in models: + success, detail = test_model_basic(model) + print_result(f"模型: {model}", success, detail) + if success: + available_models.append(model) + + print(f"\n可用模型: {len(available_models)}/{len(models)}") + + if not available_models: + print("\n❌ 没有可用的模型，测试终止") + return + + # 3. 流式输出测试 + print_separator("流式输出测试") + streaming_results = {} + for model in available_models: + result = test_streaming(model) + streaming_results[model] = result + detail = f"chunks={result['chunks']}, content_len={len(result['content'])}" + if result["error"]: + detail = f"error: {result['error']}" + print_result(f"模型: {model}", result["success"], detail) + + # 4. Thinking 模式测试 (简单任务) + print_separator("Thinking 模式测试 (简单任务)") + thinking_results = {} + for model in available_models: + result = test_thinking_mode(model, complex_task=False) + thinking_results[model] = result + detail = f"reasoning={result['has_reasoning']}, chunks={result['chunks']}" + if result["error"]: + detail = f"error: {result['error']}" + print_result(f"模型: {model}", result["success"], detail) + + # 5. Thinking 模式测试 (复杂任务) - 只测试支持 thinking 的模型 + print_separator("Thinking 模式测试 (复杂任务)") + complex_thinking_results = {} + + # 选择前3个可用模型进行复杂任务测试 + test_models = available_models[:3] + print(f"测试模型 (取前3个): {test_models}\n") + + for model in test_models: + print(f"⏳ 正在测试 {model} (复杂任务，可能需要较长时间)...") + result = test_thinking_mode(model, complex_task=True) + complex_thinking_results[model] = result + + if result["success"]: + detail = f"reasoning={result['has_reasoning']}, reasoning_len={len(result['reasoning_content'])}, content_len={len(result['content'])}" + else: + detail = f"error: {result['error']}" if result["error"] else "Unknown error" + + print_result(f"模型: {model}", result["success"], detail) + + # 如果有 reasoning 内容，打印前500字符 + if result["has_reasoning"] and result["reasoning_content"]: + print(f"\n 📝 Reasoning 内容预览 (前500字符):") + print(f" {result['reasoning_content'][:500]}...") + + # 6. 总结报告 + print_separator("测试总结报告") + + print(f"📊 模型总数: {len(models)}") + print(f"✅ 可用模型: {len(available_models)}") + print(f"❌ 不可用模型: {len(models) - len(available_models)}") + + print(f"\n📊 流式输出测试:") + streaming_pass = sum(1 for r in streaming_results.values() if r["success"]) + print(f" 通过: {streaming_pass}/{len(streaming_results)}") + + print(f"\n📊 Thinking 模式测试 (简单):") + thinking_pass = sum(1 for r in thinking_results.values() if r["success"]) + thinking_with_reasoning = sum(1 for r in thinking_results.values() if r["has_reasoning"]) + print(f" 通过: {thinking_pass}/{len(thinking_results)}") + print(f" 包含推理内容: {thinking_with_reasoning}/{len(thinking_results)}") + + print(f"\n📊 Thinking 模式测试 (复杂):") + complex_pass = sum(1 for r in complex_thinking_results.values() if r["success"]) + complex_with_reasoning = sum(1 for r in complex_thinking_results.values() if r["has_reasoning"]) + print(f" 通过: {complex_pass}/{len(complex_thinking_results)}") + print(f" 包含推理内容: {complex_with_reasoning}/{len(complex_thinking_results)}") + + # 列出所有错误 + print(f"\n📋 错误详情:") + has_errors = False + + for model, result in streaming_results.items(): + if result["error"]: + has_errors = True + print(f" [流式] {model}: {result['error'][:100]}") + + for model, result in thinking_results.items(): + if result["error"]: + has_errors = True + print(f" [Thinking简单] {model}: {result['error'][:100]}") + + for model, result in complex_thinking_results.items(): + if result["error"]: + has_errors = True + print(f" [Thinking复杂] {model}: {result['error'][:100]}") + + if not has_errors: + print(" 无错误") + + print("\n" + "="*60) + print(" 测试完成") + print("="*60 + "\n") + +def test_single_model_basic(model: str): + """单独测试一个模型的基础功能""" + print_separator(f"基础测试: {model}") + success, detail = test_model_basic(model) + print_result(f"模型: {model}", success, detail) + return success + +def test_single_model_streaming(model: str): + """单独测试一个模型的流式输出""" + print_separator(f"流式测试: {model}") + result = test_streaming(model) + detail = f"chunks={result['chunks']}, content_len={len(result['content'])}" + if result["error"]: + detail = f"error: {result['error']}" + print_result(f"模型: {model}", result["success"], detail) + if result["content"]: + print(f"\n内容: {result['content'][:300]}") + return result + +def test_single_model_thinking(model: str, complex_task: bool = False): + """单独测试一个模型的thinking模式""" + task_type = "复杂" if complex_task else "简单" + print_separator(f"Thinking测试({task_type}): {model}") + result = test_thinking_mode(model, complex_task=complex_task) + detail = f"reasoning={result['has_reasoning']}, chunks={result['chunks']}" + if result["error"]: + detail = f"error: {result['error']}" + print_result(f"模型: {model}", result["success"], detail) + if result["reasoning_content"]: + print(f"\nReasoning预览: {result['reasoning_content'][:500]}") + if result["content"]: + print(f"\n内容预览: {result['content'][:500]}") + return result + +def print_usage(): + print(""" +用法: python test_api.py [options] + +命令: + models - 获取模型列表 + basic - 测试单个模型基础功能 + stream - 测试单个模型流式输出 + thinking - 测试单个模型thinking模式(简单任务) + thinking-complex - 测试单个模型thinking模式(复杂任务) + all - 运行完整测试(原有功能) + +示例: + python test_api.py models + python test_api.py basic claude-sonnet + python test_api.py stream claude-sonnet + python test_api.py thinking claude-sonnet +""") + +if __name__ == "__main__": + import sys + + if len(sys.argv) < 2: + print_usage() + sys.exit(0) + + cmd = sys.argv[1].lower() + + if cmd == "models": + get_models() + elif cmd == "basic" and len(sys.argv) >= 3: + test_single_model_basic(sys.argv[2]) + elif cmd == "stream" and len(sys.argv) >= 3: + test_single_model_streaming(sys.argv[2]) + elif cmd == "thinking" and len(sys.argv) >= 3: + test_single_model_thinking(sys.argv[2], complex_task=False) + elif cmd == "thinking-complex" and len(sys.argv) >= 3: + test_single_model_thinking(sys.argv[2], complex_task=True) + elif cmd == "all": + run_full_test() + else: + print_usage() diff --git a/test_auth_diff.go b/test_auth_diff.go new file mode 100644 index 00000000..b294622e --- /dev/null +++ b/test_auth_diff.go @@ -0,0 +1,273 @@ +// 测试脚本 3：对比 CLIProxyAPIPlus 与官方格式的差异 +// 这个脚本分析 CLIProxyAPIPlus 保存的 token 与官方格式的差异 +// 运行方式: go run test_auth_diff.go +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +func main() { + fmt.Println("=" + strings.Repeat("=", 59)) + fmt.Println(" 测试脚本 3: Token 格式差异分析") + fmt.Println("=" + strings.Repeat("=", 59)) + + homeDir := os.Getenv("USERPROFILE") + + // 加载官方 IDE Token (Kiro IDE 生成) + fmt.Println("\n[1] 官方 Kiro IDE Token 格式") + fmt.Println("-" + strings.Repeat("-", 59)) + + ideTokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") + ideToken := loadAndAnalyze(ideTokenPath, "Kiro IDE") + + // 加载 CLIProxyAPIPlus 保存的 Token + fmt.Println("\n[2] CLIProxyAPIPlus 保存的 Token 格式") + fmt.Println("-" + strings.Repeat("-", 59)) + + cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") + files, _ := os.ReadDir(cliProxyDir) + + var cliProxyTokens []map[string]interface{} + for _, f := range files { + if strings.HasPrefix(f.Name(), "kiro") && strings.HasSuffix(f.Name(), ".json") { + p := filepath.Join(cliProxyDir, f.Name()) + token := loadAndAnalyze(p, f.Name()) + if token != nil { + cliProxyTokens = append(cliProxyTokens, token) + } + } + } + + // 对比分析 + fmt.Println("\n[3] 关键差异分析") + fmt.Println("-" + strings.Repeat("-", 59)) + + if ideToken == nil { + fmt.Println("❌ 无法加载 IDE Token，跳过对比") + } else if len(cliProxyTokens) == 0 { + fmt.Println("❌ 无法加载 CLIProxyAPIPlus Token，跳过对比") + } else { + // 对比最新的 CLIProxyAPIPlus token + cliToken := cliProxyTokens[0] + + fmt.Println("\n字段对比:") + fmt.Printf("%-20s | %-15s | %-15s\n", "字段", "IDE Token", "CLIProxy Token") + fmt.Println(strings.Repeat("-", 55)) + + fields := []string{ + "accessToken", "refreshToken", "clientId", "clientSecret", + "authMethod", "auth_method", "provider", "region", "expiresAt", "expires_at", + } + + for _, field := range fields { + ideVal := getFieldStatus(ideToken, field) + cliVal := getFieldStatus(cliToken, field) + + status := " " + if ideVal != cliVal { + if ideVal == "✅ 有" && cliVal == "❌ 无" { + status = "⚠️" + } else if ideVal == "❌ 无" && cliVal == "✅ 有" { + status = "📝" + } + } + + fmt.Printf("%-20s | %-15s | %-15s %s\n", field, ideVal, cliVal, status) + } + + // 关键问题检测 + fmt.Println("\n🔍 问题检测:") + + // 检查 clientId/clientSecret + if hasField(ideToken, "clientId") && !hasField(cliToken, "clientId") { + fmt.Println(" ⚠️ 问题: CLIProxyAPIPlus 缺少 clientId 字段!") + fmt.Println(" 原因: IdC 认证刷新 token 时需要 clientId") + } + + if hasField(ideToken, "clientSecret") && !hasField(cliToken, "clientSecret") { + fmt.Println(" ⚠️ 问题: CLIProxyAPIPlus 缺少 clientSecret 字段!") + fmt.Println(" 原因: IdC 认证刷新 token 时需要 clientSecret") + } + + // 检查字段名差异 + if hasField(cliToken, "auth_method") && !hasField(cliToken, "authMethod") { + fmt.Println(" 📝 注意: CLIProxy 使用 auth_method (snake_case)") + fmt.Println(" 而官方使用 authMethod (camelCase)") + } + + if hasField(cliToken, "expires_at") && !hasField(cliToken, "expiresAt") { + fmt.Println(" 📝 注意: CLIProxy 使用 expires_at (snake_case)") + fmt.Println(" 而官方使用 expiresAt (camelCase)") + } + } + + // Step 4: 测试使用完整格式的 token + fmt.Println("\n[4] 测试完整格式 Token (带 clientId/clientSecret)") + fmt.Println("-" + strings.Repeat("-", 59)) + + if ideToken != nil { + testWithFullToken(ideToken) + } + + fmt.Println("\n" + strings.Repeat("=", 60)) + fmt.Println(" 分析完成") + fmt.Println(strings.Repeat("=", 60)) + + // 给出建议 + fmt.Println("\n💡 修复建议:") + fmt.Println(" 1. CLIProxyAPIPlus 导入 token 时需要保留 clientId 和 clientSecret") + fmt.Println(" 2. IdC 认证刷新 token 必须使用这两个字段") + fmt.Println(" 3. 检查 CLIProxyAPIPlus 的 token 导入逻辑:") + fmt.Println(" - internal/auth/kiro/aws.go LoadKiroIDEToken()") + fmt.Println(" - sdk/auth/kiro.go ImportFromKiroIDE()") +} + +func loadAndAnalyze(path, name string) map[string]interface{} { + data, err := os.ReadFile(path) + if err != nil { + fmt.Printf("❌ 无法加载 %s: %v\n", name, err) + return nil + } + + var token map[string]interface{} + if err := json.Unmarshal(data, &token); err != nil { + fmt.Printf("❌ 无法解析 %s: %v\n", name, err) + return nil + } + + fmt.Printf("📄 %s\n", path) + fmt.Printf(" 字段数: %d\n", len(token)) + + // 列出所有字段 + fmt.Printf(" 字段列表: ") + keys := make([]string, 0, len(token)) + for k := range token { + keys = append(keys, k) + } + fmt.Printf("%v\n", keys) + + return token +} + +func getFieldStatus(token map[string]interface{}, field string) string { + if token == nil { + return "N/A" + } + if v, ok := token[field]; ok && v != nil && v != "" { + return "✅ 有" + } + return "❌ 无" +} + +func hasField(token map[string]interface{}, field string) bool { + if token == nil { + return false + } + v, ok := token[field] + return ok && v != nil && v != "" +} + +func testWithFullToken(token map[string]interface{}) { + accessToken, _ := token["accessToken"].(string) + refreshToken, _ := token["refreshToken"].(string) + clientId, _ := token["clientId"].(string) + clientSecret, _ := token["clientSecret"].(string) + region, _ := token["region"].(string) + + if region == "" { + region = "us-east-1" + } + + // 测试当前 accessToken + fmt.Println("\n测试当前 accessToken...") + if testAPICall(accessToken, region) { + fmt.Println("✅ 当前 accessToken 有效") + return + } + + fmt.Println("⚠️ 当前 accessToken 无效，尝试刷新...") + + // 检查是否有完整的刷新所需字段 + if clientId == "" || clientSecret == "" { + fmt.Println("❌ 缺少 clientId 或 clientSecret，无法刷新") + fmt.Println(" 这就是问题所在！") + return + } + + // 尝试刷新 + fmt.Println("\n使用完整字段刷新 token...") + url := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", region) + + requestBody := map[string]interface{}{ + "refreshToken": refreshToken, + "clientId": clientId, + "clientSecret": clientSecret, + "grantType": "refresh_token", + } + + body, _ := json.Marshal(requestBody) + req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf("❌ 请求失败: %v\n", err) + return + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + + if resp.StatusCode == 200 { + var refreshResp map[string]interface{} + json.Unmarshal(respBody, &refreshResp) + + newAccessToken, _ := refreshResp["accessToken"].(string) + fmt.Println("✅ Token 刷新成功!") + + // 验证新 token + if testAPICall(newAccessToken, region) { + fmt.Println("✅ 新 Token 验证成功!") + fmt.Println("\n✅ 结论: 使用完整格式 (含 clientId/clientSecret) 可以正常工作") + } + } else { + fmt.Printf("❌ 刷新失败: HTTP %d\n", resp.StatusCode) + fmt.Printf(" 响应: %s\n", string(respBody)) + } +} + +func testAPICall(accessToken, region string) bool { + url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) + + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + "isEmailRequired": true, + "resourceType": "AGENTIC_REQUEST", + } + body, _ := json.Marshal(payload) + + req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + return false + } + defer resp.Body.Close() + + return resp.StatusCode == 200 +} diff --git a/test_auth_idc_go1.go b/test_auth_idc_go1.go new file mode 100644 index 00000000..55fd5829 --- /dev/null +++ b/test_auth_idc_go1.go @@ -0,0 +1,323 @@ +// 测试脚本 1：模拟 kiro2api_go1 的 IdC 认证方式 +// 这个脚本完整模拟 kiro-gateway/temp/kiro2api_go1 的认证逻辑 +// 运行方式: go run test_auth_idc_go1.go +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "math/rand" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +// 配置常量 - 来自 kiro2api_go1/config/config.go +const ( + IdcRefreshTokenURL = "https://oidc.us-east-1.amazonaws.com/token" + CodeWhispererAPIURL = "https://codewhisperer.us-east-1.amazonaws.com" +) + +// AuthConfig - 来自 kiro2api_go1/auth/config.go +type AuthConfig struct { + AuthType string `json:"auth"` + RefreshToken string `json:"refreshToken"` + ClientID string `json:"clientId,omitempty"` + ClientSecret string `json:"clientSecret,omitempty"` +} + +// IdcRefreshRequest - 来自 kiro2api_go1/types/token.go +type IdcRefreshRequest struct { + ClientId string `json:"clientId"` + ClientSecret string `json:"clientSecret"` + GrantType string `json:"grantType"` + RefreshToken string `json:"refreshToken"` +} + +// RefreshResponse - 来自 kiro2api_go1/types/token.go +type RefreshResponse struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken,omitempty"` + ExpiresIn int `json:"expiresIn"` + TokenType string `json:"tokenType,omitempty"` +} + +// Fingerprint - 简化的指纹结构 +type Fingerprint struct { + OSType string + ConnectionBehavior string + AcceptLanguage string + SecFetchMode string + AcceptEncoding string +} + +func generateFingerprint() *Fingerprint { + osTypes := []string{"darwin", "windows", "linux"} + connections := []string{"keep-alive", "close"} + languages := []string{"en-US,en;q=0.9", "zh-CN,zh;q=0.9", "en-GB,en;q=0.9"} + fetchModes := []string{"cors", "navigate", "no-cors"} + + return &Fingerprint{ + OSType: osTypes[rand.Intn(len(osTypes))], + ConnectionBehavior: connections[rand.Intn(len(connections))], + AcceptLanguage: languages[rand.Intn(len(languages))], + SecFetchMode: fetchModes[rand.Intn(len(fetchModes))], + AcceptEncoding: "gzip, deflate, br", + } +} + +func main() { + rand.Seed(time.Now().UnixNano()) + + fmt.Println("=" + strings.Repeat("=", 59)) + fmt.Println(" 测试脚本 1: kiro2api_go1 风格 IdC 认证") + fmt.Println("=" + strings.Repeat("=", 59)) + + // Step 1: 加载官方格式的 token 文件 + fmt.Println("\n[Step 1] 加载官方格式 Token 文件") + fmt.Println("-" + strings.Repeat("-", 59)) + + // 尝试从多个位置加载 + tokenPaths := []string{ + // 优先使用包含完整 clientId/clientSecret 的文件 + "E:/ai_project_2api/kiro-gateway/configs/kiro/kiro-auth-token-1768317098.json", + filepath.Join(os.Getenv("USERPROFILE"), ".aws", "sso", "cache", "kiro-auth-token.json"), + } + + var tokenData map[string]interface{} + var loadedPath string + + for _, p := range tokenPaths { + data, err := os.ReadFile(p) + if err == nil { + if err := json.Unmarshal(data, &tokenData); err == nil { + loadedPath = p + break + } + } + } + + if tokenData == nil { + fmt.Println("❌ 无法加载任何 token 文件") + return + } + + fmt.Printf("✅ 加载文件: %s\n", loadedPath) + + // 提取关键字段 + accessToken, _ := tokenData["accessToken"].(string) + refreshToken, _ := tokenData["refreshToken"].(string) + clientId, _ := tokenData["clientId"].(string) + clientSecret, _ := tokenData["clientSecret"].(string) + authMethod, _ := tokenData["authMethod"].(string) + region, _ := tokenData["region"].(string) + + if region == "" { + region = "us-east-1" + } + + fmt.Printf("\n当前 Token 信息:\n") + fmt.Printf(" AuthMethod: %s\n", authMethod) + fmt.Printf(" Region: %s\n", region) + fmt.Printf(" AccessToken: %s...\n", truncate(accessToken, 50)) + fmt.Printf(" RefreshToken: %s...\n", truncate(refreshToken, 50)) + fmt.Printf(" ClientID: %s\n", truncate(clientId, 30)) + fmt.Printf(" ClientSecret: %s...\n", truncate(clientSecret, 50)) + + // Step 2: 验证 IdC 认证所需字段 + fmt.Println("\n[Step 2] 验证 IdC 认证必需字段") + fmt.Println("-" + strings.Repeat("-", 59)) + + missingFields := []string{} + if refreshToken == "" { + missingFields = append(missingFields, "refreshToken") + } + if clientId == "" { + missingFields = append(missingFields, "clientId") + } + if clientSecret == "" { + missingFields = append(missingFields, "clientSecret") + } + + if len(missingFields) > 0 { + fmt.Printf("❌ 缺少必需字段: %v\n", missingFields) + fmt.Println(" IdC 认证需要: refreshToken, clientId, clientSecret") + return + } + fmt.Println("✅ 所有必需字段都存在") + + // Step 3: 测试直接使用 accessToken 调用 API + fmt.Println("\n[Step 3] 测试当前 AccessToken 有效性") + fmt.Println("-" + strings.Repeat("-", 59)) + + if testAPICall(accessToken, region) { + fmt.Println("✅ 当前 AccessToken 有效，无需刷新") + } else { + fmt.Println("⚠️ 当前 AccessToken 无效，需要刷新") + + // Step 4: 使用 kiro2api_go1 风格刷新 token + fmt.Println("\n[Step 4] 使用 kiro2api_go1 风格刷新 Token") + fmt.Println("-" + strings.Repeat("-", 59)) + + newToken, err := refreshIdCToken(AuthConfig{ + AuthType: "IdC", + RefreshToken: refreshToken, + ClientID: clientId, + ClientSecret: clientSecret, + }, region) + + if err != nil { + fmt.Printf("❌ 刷新失败: %v\n", err) + return + } + + fmt.Println("✅ Token 刷新成功！") + fmt.Printf(" 新 AccessToken: %s...\n", truncate(newToken.AccessToken, 50)) + fmt.Printf(" ExpiresIn: %d 秒\n", newToken.ExpiresIn) + + // Step 5: 验证新 token + fmt.Println("\n[Step 5] 验证新 Token") + fmt.Println("-" + strings.Repeat("-", 59)) + + if testAPICall(newToken.AccessToken, region) { + fmt.Println("✅ 新 Token 验证成功！") + + // 保存新 token + saveNewToken(loadedPath, newToken, tokenData) + } else { + fmt.Println("❌ 新 Token 验证失败") + } + } + + fmt.Println("\n" + strings.Repeat("=", 60)) + fmt.Println(" 测试完成") + fmt.Println(strings.Repeat("=", 60)) +} + +// refreshIdCToken - 完全模拟 kiro2api_go1/auth/refresh.go 的 refreshIdCToken 函数 +func refreshIdCToken(authConfig AuthConfig, region string) (*RefreshResponse, error) { + refreshReq := IdcRefreshRequest{ + ClientId: authConfig.ClientID, + ClientSecret: authConfig.ClientSecret, + GrantType: "refresh_token", + RefreshToken: authConfig.RefreshToken, + } + + reqBody, err := json.Marshal(refreshReq) + if err != nil { + return nil, fmt.Errorf("序列化IdC请求失败: %v", err) + } + + url := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", region) + req, err := http.NewRequest("POST", url, bytes.NewBuffer(reqBody)) + if err != nil { + return nil, fmt.Errorf("创建IdC请求失败: %v", err) + } + + // 设置 IdC 特殊 headers（使用指纹随机化）- 完全模拟 kiro2api_go1 + fp := generateFingerprint() + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region)) + req.Header.Set("Connection", fp.ConnectionBehavior) + req.Header.Set("x-amz-user-agent", fmt.Sprintf("aws-sdk-js/3.738.0 ua/2.1 os/%s lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE", fp.OSType)) + req.Header.Set("Accept", "*/*") + req.Header.Set("Accept-Language", fp.AcceptLanguage) + req.Header.Set("sec-fetch-mode", fp.SecFetchMode) + req.Header.Set("User-Agent", "node") + req.Header.Set("Accept-Encoding", fp.AcceptEncoding) + + fmt.Println("发送刷新请求:") + fmt.Printf(" URL: %s\n", url) + fmt.Println(" Headers:") + for k, v := range req.Header { + if k == "Content-Type" || k == "Host" || k == "X-Amz-User-Agent" || k == "User-Agent" { + fmt.Printf(" %s: %s\n", k, v[0]) + } + } + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("IdC请求失败: %v", err) + } + defer resp.Body.Close() + + body, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("IdC刷新失败: 状态码 %d, 响应: %s", resp.StatusCode, string(body)) + } + + var refreshResp RefreshResponse + if err := json.Unmarshal(body, &refreshResp); err != nil { + return nil, fmt.Errorf("解析IdC响应失败: %v", err) + } + + return &refreshResp, nil +} + +func testAPICall(accessToken, region string) bool { + url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) + + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + "isEmailRequired": true, + "resourceType": "AGENTIC_REQUEST", + } + body, _ := json.Marshal(payload) + + req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf(" 请求错误: %v\n", err) + return false + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + fmt.Printf(" API 响应: HTTP %d\n", resp.StatusCode) + + if resp.StatusCode == 200 { + return true + } + + fmt.Printf(" 错误详情: %s\n", truncate(string(respBody), 200)) + return false +} + +func saveNewToken(originalPath string, newToken *RefreshResponse, originalData map[string]interface{}) { + // 更新 token 数据 + originalData["accessToken"] = newToken.AccessToken + if newToken.RefreshToken != "" { + originalData["refreshToken"] = newToken.RefreshToken + } + originalData["expiresAt"] = time.Now().Add(time.Duration(newToken.ExpiresIn) * time.Second).Format(time.RFC3339) + + data, _ := json.MarshalIndent(originalData, "", " ") + + // 保存到新文件 + newPath := strings.TrimSuffix(originalPath, ".json") + "_refreshed.json" + if err := os.WriteFile(newPath, data, 0644); err != nil { + fmt.Printf("⚠️ 保存失败: %v\n", err) + } else { + fmt.Printf("✅ 新 Token 已保存到: %s\n", newPath) + } +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] +} diff --git a/test_auth_js_style.go b/test_auth_js_style.go new file mode 100644 index 00000000..6ded3305 --- /dev/null +++ b/test_auth_js_style.go @@ -0,0 +1,237 @@ +// 测试脚本 2：模拟 kiro2Api_js 的认证方式 +// 这个脚本完整模拟 kiro-gateway/temp/kiro2Api_js 的认证逻辑 +// 运行方式: go run test_auth_js_style.go +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +// 常量 - 来自 kiro2Api_js/src/kiro/auth.js +const ( + REFRESH_URL_TEMPLATE = "https://prod.{{region}}.auth.desktop.kiro.dev/refreshToken" + REFRESH_IDC_URL_TEMPLATE = "https://oidc.{{region}}.amazonaws.com/token" + AUTH_METHOD_SOCIAL = "social" + AUTH_METHOD_IDC = "IdC" +) + +func main() { + fmt.Println("=" + strings.Repeat("=", 59)) + fmt.Println(" 测试脚本 2: kiro2Api_js 风格认证") + fmt.Println("=" + strings.Repeat("=", 59)) + + // Step 1: 加载 token 文件 + fmt.Println("\n[Step 1] 加载 Token 文件") + fmt.Println("-" + strings.Repeat("-", 59)) + + tokenPaths := []string{ + filepath.Join(os.Getenv("USERPROFILE"), ".aws", "sso", "cache", "kiro-auth-token.json"), + "E:/ai_project_2api/kiro-gateway/configs/kiro/kiro-auth-token-1768317098.json", + } + + var tokenData map[string]interface{} + var loadedPath string + + for _, p := range tokenPaths { + data, err := os.ReadFile(p) + if err == nil { + if err := json.Unmarshal(data, &tokenData); err == nil { + loadedPath = p + break + } + } + } + + if tokenData == nil { + fmt.Println("❌ 无法加载任何 token 文件") + return + } + + fmt.Printf("✅ 加载文件: %s\n", loadedPath) + + // 提取字段 - 模拟 kiro2Api_js/src/kiro/auth.js initializeAuth + accessToken, _ := tokenData["accessToken"].(string) + refreshToken, _ := tokenData["refreshToken"].(string) + clientId, _ := tokenData["clientId"].(string) + clientSecret, _ := tokenData["clientSecret"].(string) + authMethod, _ := tokenData["authMethod"].(string) + region, _ := tokenData["region"].(string) + + if region == "" { + region = "us-east-1" + fmt.Println("⚠️ Region 未设置，使用默认值 us-east-1") + } + + fmt.Printf("\nToken 信息:\n") + fmt.Printf(" AuthMethod: %s\n", authMethod) + fmt.Printf(" Region: %s\n", region) + fmt.Printf(" 有 ClientID: %v\n", clientId != "") + fmt.Printf(" 有 ClientSecret: %v\n", clientSecret != "") + + // Step 2: 测试当前 token + fmt.Println("\n[Step 2] 测试当前 AccessToken") + fmt.Println("-" + strings.Repeat("-", 59)) + + if testAPI(accessToken, region) { + fmt.Println("✅ 当前 AccessToken 有效") + return + } + + fmt.Println("⚠️ 当前 AccessToken 无效，开始刷新...") + + // Step 3: 根据 authMethod 选择刷新方式 - 模拟 doRefreshToken + fmt.Println("\n[Step 3] 刷新 Token (JS 风格)") + fmt.Println("-" + strings.Repeat("-", 59)) + + var refreshURL string + var requestBody map[string]interface{} + + // 判断认证方式 - 模拟 kiro2Api_js auth.js doRefreshToken + if authMethod == AUTH_METHOD_SOCIAL { + // Social 认证 + refreshURL = strings.Replace(REFRESH_URL_TEMPLATE, "{{region}}", region, 1) + requestBody = map[string]interface{}{ + "refreshToken": refreshToken, + } + fmt.Println("使用 Social 认证方式") + } else { + // IdC 认证 (默认) + refreshURL = strings.Replace(REFRESH_IDC_URL_TEMPLATE, "{{region}}", region, 1) + requestBody = map[string]interface{}{ + "refreshToken": refreshToken, + "clientId": clientId, + "clientSecret": clientSecret, + "grantType": "refresh_token", + } + fmt.Println("使用 IdC 认证方式") + } + + fmt.Printf("刷新 URL: %s\n", refreshURL) + fmt.Printf("请求字段: %v\n", getKeys(requestBody)) + + // 发送刷新请求 + body, _ := json.Marshal(requestBody) + req, _ := http.NewRequest("POST", refreshURL, bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf("❌ 请求失败: %v\n", err) + return + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + + fmt.Printf("\n响应状态: HTTP %d\n", resp.StatusCode) + + if resp.StatusCode != 200 { + fmt.Printf("❌ 刷新失败: %s\n", string(respBody)) + + // 分析错误 + var errResp map[string]interface{} + if err := json.Unmarshal(respBody, &errResp); err == nil { + if errType, ok := errResp["error"].(string); ok { + fmt.Printf("错误类型: %s\n", errType) + if errType == "invalid_grant" { + fmt.Println("\n💡 提示: refresh_token 可能已过期，需要重新授权") + } + } + if errDesc, ok := errResp["error_description"].(string); ok { + fmt.Printf("错误描述: %s\n", errDesc) + } + } + return + } + + // 解析响应 + var refreshResp map[string]interface{} + json.Unmarshal(respBody, &refreshResp) + + newAccessToken, _ := refreshResp["accessToken"].(string) + newRefreshToken, _ := refreshResp["refreshToken"].(string) + expiresIn, _ := refreshResp["expiresIn"].(float64) + + fmt.Println("✅ Token 刷新成功!") + fmt.Printf(" 新 AccessToken: %s...\n", truncate(newAccessToken, 50)) + fmt.Printf(" ExpiresIn: %.0f 秒\n", expiresIn) + if newRefreshToken != "" { + fmt.Printf(" 新 RefreshToken: %s...\n", truncate(newRefreshToken, 50)) + } + + // Step 4: 验证新 token + fmt.Println("\n[Step 4] 验证新 Token") + fmt.Println("-" + strings.Repeat("-", 59)) + + if testAPI(newAccessToken, region) { + fmt.Println("✅ 新 Token 验证成功!") + + // 保存新 token - 模拟 saveCredentialsToFile + tokenData["accessToken"] = newAccessToken + if newRefreshToken != "" { + tokenData["refreshToken"] = newRefreshToken + } + tokenData["expiresAt"] = time.Now().Add(time.Duration(expiresIn) * time.Second).Format(time.RFC3339) + + saveData, _ := json.MarshalIndent(tokenData, "", " ") + newPath := strings.TrimSuffix(loadedPath, ".json") + "_js_refreshed.json" + os.WriteFile(newPath, saveData, 0644) + fmt.Printf("✅ 已保存到: %s\n", newPath) + } else { + fmt.Println("❌ 新 Token 验证失败") + } + + fmt.Println("\n" + strings.Repeat("=", 60)) + fmt.Println(" 测试完成") + fmt.Println(strings.Repeat("=", 60)) +} + +func testAPI(accessToken, region string) bool { + url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) + + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + "isEmailRequired": true, + "resourceType": "AGENTIC_REQUEST", + } + body, _ := json.Marshal(payload) + + req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + return false + } + defer resp.Body.Close() + + return resp.StatusCode == 200 +} + +func getKeys(m map[string]interface{}) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + return keys +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] +} diff --git a/test_kiro_debug.go b/test_kiro_debug.go new file mode 100644 index 00000000..0fbbed6c --- /dev/null +++ b/test_kiro_debug.go @@ -0,0 +1,348 @@ +// 独立测试脚本：排查 Kiro Token 403 错误 +// 运行方式: go run test_kiro_debug.go +package main + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +// Token 结构 - 匹配 Kiro IDE 格式 +type KiroIDEToken struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken"` + ExpiresAt string `json:"expiresAt"` + ClientIDHash string `json:"clientIdHash,omitempty"` + AuthMethod string `json:"authMethod"` + Provider string `json:"provider"` + Region string `json:"region,omitempty"` +} + +// Token 结构 - 匹配 CLIProxyAPIPlus 格式 +type CLIProxyToken struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + ProfileArn string `json:"profile_arn"` + ExpiresAt string `json:"expires_at"` + AuthMethod string `json:"auth_method"` + Provider string `json:"provider"` + ClientID string `json:"client_id,omitempty"` + ClientSecret string `json:"client_secret,omitempty"` + Email string `json:"email,omitempty"` + Type string `json:"type"` +} + +func main() { + fmt.Println("=" + strings.Repeat("=", 59)) + fmt.Println(" Kiro Token 403 错误排查工具") + fmt.Println("=" + strings.Repeat("=", 59)) + + homeDir, _ := os.UserHomeDir() + + // Step 1: 检查 Kiro IDE Token 文件 + fmt.Println("\n[Step 1] 检查 Kiro IDE Token 文件") + fmt.Println("-" + strings.Repeat("-", 59)) + + ideTokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") + ideToken, err := loadKiroIDEToken(ideTokenPath) + if err != nil { + fmt.Printf("❌ 无法加载 Kiro IDE Token: %v\n", err) + return + } + fmt.Printf("✅ Token 文件: %s\n", ideTokenPath) + fmt.Printf(" AuthMethod: %s\n", ideToken.AuthMethod) + fmt.Printf(" Provider: %s\n", ideToken.Provider) + fmt.Printf(" Region: %s\n", ideToken.Region) + fmt.Printf(" ExpiresAt: %s\n", ideToken.ExpiresAt) + fmt.Printf(" AccessToken (前50字符): %s...\n", truncate(ideToken.AccessToken, 50)) + + // Step 2: 检查 Token 过期状态 + fmt.Println("\n[Step 2] 检查 Token 过期状态") + fmt.Println("-" + strings.Repeat("-", 59)) + + expiresAt, err := parseExpiresAt(ideToken.ExpiresAt) + if err != nil { + fmt.Printf("❌ 无法解析过期时间: %v\n", err) + } else { + now := time.Now() + if now.After(expiresAt) { + fmt.Printf("❌ Token 已过期！过期时间: %s，当前时间: %s\n", expiresAt.Format(time.RFC3339), now.Format(time.RFC3339)) + } else { + remaining := expiresAt.Sub(now) + fmt.Printf("✅ Token 未过期，剩余: %s\n", remaining.Round(time.Second)) + } + } + + // Step 3: 检查 CLIProxyAPIPlus 保存的 Token + fmt.Println("\n[Step 3] 检查 CLIProxyAPIPlus 保存的 Token") + fmt.Println("-" + strings.Repeat("-", 59)) + + cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") + files, _ := os.ReadDir(cliProxyDir) + for _, f := range files { + if strings.HasPrefix(f.Name(), "kiro") && strings.HasSuffix(f.Name(), ".json") { + filePath := filepath.Join(cliProxyDir, f.Name()) + cliToken, err := loadCLIProxyToken(filePath) + if err != nil { + fmt.Printf("❌ %s: 加载失败 - %v\n", f.Name(), err) + continue + } + fmt.Printf("📄 %s:\n", f.Name()) + fmt.Printf(" AuthMethod: %s\n", cliToken.AuthMethod) + fmt.Printf(" Provider: %s\n", cliToken.Provider) + fmt.Printf(" ExpiresAt: %s\n", cliToken.ExpiresAt) + fmt.Printf(" AccessToken (前50字符): %s...\n", truncate(cliToken.AccessToken, 50)) + + // 比较 Token + if cliToken.AccessToken == ideToken.AccessToken { + fmt.Printf(" ✅ AccessToken 与 IDE Token 一致\n") + } else { + fmt.Printf(" ⚠️ AccessToken 与 IDE Token 不一致！\n") + } + } + } + + // Step 4: 直接测试 Token 有效性 (调用 Kiro API) + fmt.Println("\n[Step 4] 直接测试 Token 有效性") + fmt.Println("-" + strings.Repeat("-", 59)) + + testTokenValidity(ideToken.AccessToken, ideToken.Region) + + // Step 5: 测试不同的请求头格式 + fmt.Println("\n[Step 5] 测试不同的请求头格式") + fmt.Println("-" + strings.Repeat("-", 59)) + + testDifferentHeaders(ideToken.AccessToken, ideToken.Region) + + // Step 6: 解析 JWT 内容 + fmt.Println("\n[Step 6] 解析 JWT Token 内容") + fmt.Println("-" + strings.Repeat("-", 59)) + + parseJWT(ideToken.AccessToken) + + fmt.Println("\n" + strings.Repeat("=", 60)) + fmt.Println(" 排查完成") + fmt.Println(strings.Repeat("=", 60)) +} + +func loadKiroIDEToken(path string) (*KiroIDEToken, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var token KiroIDEToken + if err := json.Unmarshal(data, &token); err != nil { + return nil, err + } + return &token, nil +} + +func loadCLIProxyToken(path string) (*CLIProxyToken, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var token CLIProxyToken + if err := json.Unmarshal(data, &token); err != nil { + return nil, err + } + return &token, nil +} + +func parseExpiresAt(s string) (time.Time, error) { + formats := []string{ + time.RFC3339, + "2006-01-02T15:04:05.000Z", + "2006-01-02T15:04:05Z", + } + for _, f := range formats { + if t, err := time.Parse(f, s); err == nil { + return t, nil + } + } + return time.Time{}, fmt.Errorf("无法解析时间格式: %s", s) +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] +} + +func testTokenValidity(accessToken, region string) { + if region == "" { + region = "us-east-1" + } + + // 测试 GetUsageLimits API + url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) + + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + "isEmailRequired": true, + "resourceType": "AGENTIC_REQUEST", + } + body, _ := json.Marshal(payload) + + req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + fmt.Printf("请求 URL: %s\n", url) + fmt.Printf("请求头:\n") + for k, v := range req.Header { + if k == "Authorization" { + fmt.Printf(" %s: Bearer %s...\n", k, truncate(v[0][7:], 30)) + } else { + fmt.Printf(" %s: %s\n", k, v[0]) + } + } + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf("❌ 请求失败: %v\n", err) + return + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + fmt.Printf("响应状态: %d\n", resp.StatusCode) + fmt.Printf("响应内容: %s\n", string(respBody)) + + if resp.StatusCode == 200 { + fmt.Println("✅ Token 有效！") + } else if resp.StatusCode == 403 { + fmt.Println("❌ Token 无效或已过期 (403)") + } +} + +func testDifferentHeaders(accessToken, region string) { + if region == "" { + region = "us-east-1" + } + + tests := []struct { + name string + headers map[string]string + }{ + { + name: "最小请求头", + headers: map[string]string{ + "Content-Type": "application/json", + "Authorization": "Bearer " + accessToken, + }, + }, + { + name: "模拟 kiro2api_go1 风格", + headers: map[string]string{ + "Content-Type": "application/json", + "Accept": "text/event-stream", + "Authorization": "Bearer " + accessToken, + "x-amzn-kiro-agent-mode": "vibe", + "x-amzn-codewhisperer-optout": "true", + "amz-sdk-invocation-id": "test-invocation-id", + "amz-sdk-request": "attempt=1; max=3", + "x-amz-user-agent": "aws-sdk-js/1.0.27 KiroIDE-0.8.0-abc123", + "User-Agent": "aws-sdk-js/1.0.27 ua/2.1 os/windows#10.0 lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.27 m/E KiroIDE-0.8.0-abc123", + }, + }, + { + name: "模拟 CLIProxyAPIPlus 风格", + headers: map[string]string{ + "Content-Type": "application/x-amz-json-1.0", + "x-amz-target": "AmazonCodeWhispererService.GetUsageLimits", + "Authorization": "Bearer " + accessToken, + "Accept": "application/json", + "amz-sdk-invocation-id": "test-invocation-id", + "amz-sdk-request": "attempt=1; max=1", + "Connection": "close", + }, + }, + } + + url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + "isEmailRequired": true, + "resourceType": "AGENTIC_REQUEST", + } + body, _ := json.Marshal(payload) + + for _, test := range tests { + fmt.Printf("\n测试: %s\n", test.name) + + req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) + for k, v := range test.headers { + req.Header.Set(k, v) + } + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf(" ❌ 请求失败: %v\n", err) + continue + } + + respBody, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + if resp.StatusCode == 200 { + fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) + } else { + fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) + } + } +} + +func parseJWT(token string) { + parts := strings.Split(token, ".") + if len(parts) < 2 { + fmt.Println("Token 不是 JWT 格式") + return + } + + // 解码 header + headerData, err := base64.RawURLEncoding.DecodeString(parts[0]) + if err != nil { + fmt.Printf("无法解码 JWT header: %v\n", err) + } else { + var header map[string]interface{} + json.Unmarshal(headerData, &header) + fmt.Printf("JWT Header: %v\n", header) + } + + // 解码 payload + payloadData, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + fmt.Printf("无法解码 JWT payload: %v\n", err) + } else { + var payload map[string]interface{} + json.Unmarshal(payloadData, &payload) + fmt.Printf("JWT Payload:\n") + for k, v := range payload { + fmt.Printf(" %s: %v\n", k, v) + } + + // 检查过期时间 + if exp, ok := payload["exp"].(float64); ok { + expTime := time.Unix(int64(exp), 0) + if time.Now().After(expTime) { + fmt.Printf(" ⚠️ JWT 已过期! exp=%s\n", expTime.Format(time.RFC3339)) + } else { + fmt.Printf(" ✅ JWT 未过期, 剩余: %s\n", expTime.Sub(time.Now()).Round(time.Second)) + } + } + } +} diff --git a/test_proxy_debug.go b/test_proxy_debug.go new file mode 100644 index 00000000..82369e74 --- /dev/null +++ b/test_proxy_debug.go @@ -0,0 +1,367 @@ +// 测试脚本 2：通过 CLIProxyAPIPlus 代理层排查问题 +// 运行方式: go run test_proxy_debug.go +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +const ( + ProxyURL = "http://localhost:8317" + APIKey = "your-api-key-1" +) + +func main() { + fmt.Println("=" + strings.Repeat("=", 59)) + fmt.Println(" CLIProxyAPIPlus 代理层问题排查") + fmt.Println("=" + strings.Repeat("=", 59)) + + // Step 1: 检查代理服务状态 + fmt.Println("\n[Step 1] 检查代理服务状态") + fmt.Println("-" + strings.Repeat("-", 59)) + + resp, err := http.Get(ProxyURL + "/health") + if err != nil { + fmt.Printf("❌ 代理服务不可达: %v\n", err) + fmt.Println("请确保服务正在运行: go run ./cmd/server/main.go") + return + } + resp.Body.Close() + fmt.Printf("✅ 代理服务正常 (HTTP %d)\n", resp.StatusCode) + + // Step 2: 获取模型列表 + fmt.Println("\n[Step 2] 获取模型列表") + fmt.Println("-" + strings.Repeat("-", 59)) + + models := getModels() + if len(models) == 0 { + fmt.Println("❌ 没有可用的模型，检查凭据加载") + checkCredentials() + return + } + fmt.Printf("✅ 找到 %d 个模型:\n", len(models)) + for _, m := range models { + fmt.Printf(" - %s\n", m) + } + + // Step 3: 测试模型请求 - 捕获详细错误 + fmt.Println("\n[Step 3] 测试模型请求（详细日志）") + fmt.Println("-" + strings.Repeat("-", 59)) + + if len(models) > 0 { + testModel := models[0] + testModelRequest(testModel) + } + + // Step 4: 检查代理内部 Token 状态 + fmt.Println("\n[Step 4] 检查代理服务加载的凭据") + fmt.Println("-" + strings.Repeat("-", 59)) + + checkProxyCredentials() + + // Step 5: 对比直接请求和代理请求 + fmt.Println("\n[Step 5] 对比直接请求 vs 代理请求") + fmt.Println("-" + strings.Repeat("-", 59)) + + compareDirectVsProxy() + + fmt.Println("\n" + strings.Repeat("=", 60)) + fmt.Println(" 排查完成") + fmt.Println(strings.Repeat("=", 60)) +} + +func getModels() []string { + req, _ := http.NewRequest("GET", ProxyURL+"/v1/models", nil) + req.Header.Set("Authorization", "Bearer "+APIKey) + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf("❌ 请求失败: %v\n", err) + return nil + } + defer resp.Body.Close() + + body, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != 200 { + fmt.Printf("❌ HTTP %d: %s\n", resp.StatusCode, string(body)) + return nil + } + + var result struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + json.Unmarshal(body, &result) + + models := make([]string, len(result.Data)) + for i, m := range result.Data { + models[i] = m.ID + } + return models +} + +func checkCredentials() { + homeDir, _ := os.UserHomeDir() + cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") + + fmt.Printf("\n检查凭据目录: %s\n", cliProxyDir) + files, err := os.ReadDir(cliProxyDir) + if err != nil { + fmt.Printf("❌ 无法读取目录: %v\n", err) + return + } + + for _, f := range files { + if strings.HasSuffix(f.Name(), ".json") { + fmt.Printf(" 📄 %s\n", f.Name()) + } + } +} + +func testModelRequest(model string) { + fmt.Printf("测试模型: %s\n", model) + + payload := map[string]interface{}{ + "model": model, + "messages": []map[string]string{ + {"role": "user", "content": "Say 'OK' if you receive this."}, + }, + "max_tokens": 50, + "stream": false, + } + body, _ := json.Marshal(payload) + + req, _ := http.NewRequest("POST", ProxyURL+"/v1/chat/completions", bytes.NewBuffer(body)) + req.Header.Set("Authorization", "Bearer "+APIKey) + req.Header.Set("Content-Type", "application/json") + + fmt.Println("\n发送请求:") + fmt.Printf(" URL: %s/v1/chat/completions\n", ProxyURL) + fmt.Printf(" Model: %s\n", model) + + client := &http.Client{Timeout: 60 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf("❌ 请求失败: %v\n", err) + return + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + + fmt.Printf("\n响应:\n") + fmt.Printf(" Status: %d\n", resp.StatusCode) + fmt.Printf(" Headers:\n") + for k, v := range resp.Header { + fmt.Printf(" %s: %s\n", k, strings.Join(v, ", ")) + } + + // 格式化 JSON 输出 + var prettyJSON bytes.Buffer + if err := json.Indent(&prettyJSON, respBody, " ", " "); err == nil { + fmt.Printf(" Body:\n %s\n", prettyJSON.String()) + } else { + fmt.Printf(" Body: %s\n", string(respBody)) + } + + if resp.StatusCode == 200 { + fmt.Println("\n✅ 请求成功！") + } else { + fmt.Println("\n❌ 请求失败！分析错误原因...") + analyzeError(respBody) + } +} + +func analyzeError(body []byte) { + var errResp struct { + Message string `json:"message"` + Reason string `json:"reason"` + Error struct { + Message string `json:"message"` + Type string `json:"type"` + } `json:"error"` + } + json.Unmarshal(body, &errResp) + + if errResp.Message != "" { + fmt.Printf("错误消息: %s\n", errResp.Message) + } + if errResp.Reason != "" { + fmt.Printf("错误原因: %s\n", errResp.Reason) + } + if errResp.Error.Message != "" { + fmt.Printf("错误详情: %s (类型: %s)\n", errResp.Error.Message, errResp.Error.Type) + } + + // 分析常见错误 + bodyStr := string(body) + if strings.Contains(bodyStr, "bearer token") || strings.Contains(bodyStr, "invalid") { + fmt.Println("\n可能的原因:") + fmt.Println(" 1. Token 已过期 - 需要刷新") + fmt.Println(" 2. Token 格式不正确 - 检查凭据文件") + fmt.Println(" 3. 代理服务加载了旧的 Token") + } +} + +func checkProxyCredentials() { + // 尝试通过管理 API 获取凭据状态 + req, _ := http.NewRequest("GET", ProxyURL+"/v0/management/auth/list", nil) + // 使用配置中的管理密钥 admin123 + req.Header.Set("Authorization", "Bearer admin123") + + client := &http.Client{Timeout: 10 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf("❌ 无法访问管理 API: %v\n", err) + return + } + defer resp.Body.Close() + + body, _ := io.ReadAll(resp.Body) + + if resp.StatusCode == 200 { + fmt.Println("管理 API 返回的凭据列表:") + var prettyJSON bytes.Buffer + if err := json.Indent(&prettyJSON, body, " ", " "); err == nil { + fmt.Printf("%s\n", prettyJSON.String()) + } else { + fmt.Printf("%s\n", string(body)) + } + } else { + fmt.Printf("管理 API 返回: HTTP %d\n", resp.StatusCode) + fmt.Printf("响应: %s\n", truncate(string(body), 200)) + } +} + +func compareDirectVsProxy() { + homeDir, _ := os.UserHomeDir() + tokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") + + data, err := os.ReadFile(tokenPath) + if err != nil { + fmt.Printf("❌ 无法读取 Token 文件: %v\n", err) + return + } + + var token struct { + AccessToken string `json:"accessToken"` + Region string `json:"region"` + } + json.Unmarshal(data, &token) + + if token.Region == "" { + token.Region = "us-east-1" + } + + // 直接请求 + fmt.Println("\n1. 直接请求 Kiro API:") + directSuccess := testDirectKiroAPI(token.AccessToken, token.Region) + + // 通过代理请求 + fmt.Println("\n2. 通过代理请求:") + proxySuccess := testProxyAPI() + + // 结论 + fmt.Println("\n结论:") + if directSuccess && !proxySuccess { + fmt.Println(" ⚠️ 直接请求成功，代理请求失败") + fmt.Println(" 问题在于 CLIProxyAPIPlus 代理层") + fmt.Println(" 可能原因:") + fmt.Println(" 1. 代理服务使用了过期的 Token") + fmt.Println(" 2. Token 刷新逻辑有问题") + fmt.Println(" 3. 请求头构造不正确") + } else if directSuccess && proxySuccess { + fmt.Println(" ✅ 两者都成功") + } else if !directSuccess && !proxySuccess { + fmt.Println(" ❌ 两者都失败 - Token 本身可能有问题") + } +} + +func testDirectKiroAPI(accessToken, region string) bool { + url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) + + payload := map[string]interface{}{ + "origin": "AI_EDITOR", + "isEmailRequired": true, + "resourceType": "AGENTIC_REQUEST", + } + body, _ := json.Marshal(payload) + + req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/x-amz-json-1.0") + req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf(" ❌ 请求失败: %v\n", err) + return false + } + defer resp.Body.Close() + + if resp.StatusCode == 200 { + fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) + return true + } + respBody, _ := io.ReadAll(resp.Body) + fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) + return false +} + +func testProxyAPI() bool { + models := getModels() + if len(models) == 0 { + fmt.Println(" ❌ 没有可用模型") + return false + } + + payload := map[string]interface{}{ + "model": models[0], + "messages": []map[string]string{ + {"role": "user", "content": "Say OK"}, + }, + "max_tokens": 10, + "stream": false, + } + body, _ := json.Marshal(payload) + + req, _ := http.NewRequest("POST", ProxyURL+"/v1/chat/completions", bytes.NewBuffer(body)) + req.Header.Set("Authorization", "Bearer "+APIKey) + req.Header.Set("Content-Type", "application/json") + + client := &http.Client{Timeout: 60 * time.Second} + resp, err := client.Do(req) + if err != nil { + fmt.Printf(" ❌ 请求失败: %v\n", err) + return false + } + defer resp.Body.Close() + + if resp.StatusCode == 200 { + fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) + return true + } + respBody, _ := io.ReadAll(resp.Body) + fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) + return false +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "..." +} From c9301a6d18d676f42e7b6b6d5e3084730c9cd52a Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Sun, 18 Jan 2026 15:07:29 +0800 Subject: [PATCH 02/10] docs: update README with new features and Docker deployment guide --- README.md | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++++ README_CN.md | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/README.md b/README.md index d00e91c9..1555e643 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,67 @@ The Plus release stays in lockstep with the mainline features. - Added GitHub Copilot support (OAuth login), provided by [em4go](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth) - Added Kiro (AWS CodeWhisperer) support (OAuth login), provided by [fuko2935](https://github.com/fuko2935/CLIProxyAPI/tree/feature/kiro-integration), [Ravens2121](https://github.com/Ravens2121/CLIProxyAPIPlus/) +## New Features (Plus Enhanced) + +- **OAuth Web Authentication**: Browser-based OAuth login for Kiro with beautiful web UI +- **Rate Limiter**: Built-in request rate limiting to prevent API abuse +- **Background Token Refresh**: Automatic token refresh in background to avoid expiration +- **Metrics & Monitoring**: Request metrics collection for monitoring and debugging +- **Device Fingerprint**: Device fingerprint generation for enhanced security +- **Cooldown Management**: Smart cooldown mechanism for API rate limits +- **Usage Checker**: Real-time usage monitoring and quota management +- **Model Converter**: Unified model name conversion across providers +- **UTF-8 Stream Processing**: Improved streaming response handling + +## Quick Deployment with Docker + +### One-Command Deployment + +```bash +# Create deployment directory +mkdir -p ~/cli-proxy && cd ~/cli-proxy + +# Create docker-compose.yml +cat > docker-compose.yml << 'EOF' +services: + cli-proxy-api: + image: 17600006524/cli-proxy-api-plus:latest + container_name: cli-proxy-api-plus + ports: + - "8317:8317" + volumes: + - ./config.yaml:/CLIProxyAPI/config.yaml + - ./auths:/root/.cli-proxy-api + - ./logs:/CLIProxyAPI/logs + restart: unless-stopped +EOF + +# Download example config +curl -o config.yaml https://raw.githubusercontent.com/linlang781/CLIProxyAPIPlus/main/config.example.yaml + +# Pull and start +docker compose pull && docker compose up -d +``` + +### Configuration + +Edit `config.yaml` before starting: + +```yaml +# Basic configuration example +server: + port: 8317 + +# Add your provider configurations here +``` + +### Update to Latest Version + +```bash +cd ~/cli-proxy +docker compose pull && docker compose up -d +``` + ## Contributing This project only accepts pull requests that relate to third-party provider support. Any pull requests unrelated to third-party provider support will be rejected. diff --git a/README_CN.md b/README_CN.md index 21132b86..6ac2e483 100644 --- a/README_CN.md +++ b/README_CN.md @@ -13,6 +13,67 @@ - 新增 GitHub Copilot 支持（OAuth 登录），由[em4go](https://github.com/em4go/CLIProxyAPI/tree/feature/github-copilot-auth)提供 - 新增 Kiro (AWS CodeWhisperer) 支持 (OAuth 登录), 由[fuko2935](https://github.com/fuko2935/CLIProxyAPI/tree/feature/kiro-integration)、[Ravens2121](https://github.com/Ravens2121/CLIProxyAPIPlus/)提供 +## 新增功能 (Plus 增强版) + +- **OAuth Web 认证**: 基于浏览器的 Kiro OAuth 登录，提供美观的 Web UI +- **请求限流器**: 内置请求限流，防止 API 滥用 +- **后台令牌刷新**: 自动后台刷新令牌，避免过期 +- **监控指标**: 请求指标收集，用于监控和调试 +- **设备指纹**: 设备指纹生成，增强安全性 +- **冷却管理**: 智能冷却机制，应对 API 速率限制 +- **用量检查器**: 实时用量监控和配额管理 +- **模型转换器**: 跨供应商的统一模型名称转换 +- **UTF-8 流处理**: 改进的流式响应处理 + +## Docker 快速部署 + +### 一键部署 + +```bash +# 创建部署目录 +mkdir -p ~/cli-proxy && cd ~/cli-proxy + +# 创建 docker-compose.yml +cat > docker-compose.yml << 'EOF' +services: + cli-proxy-api: + image: 17600006524/cli-proxy-api-plus:latest + container_name: cli-proxy-api-plus + ports: + - "8317:8317" + volumes: + - ./config.yaml:/CLIProxyAPI/config.yaml + - ./auths:/root/.cli-proxy-api + - ./logs:/CLIProxyAPI/logs + restart: unless-stopped +EOF + +# 下载示例配置 +curl -o config.yaml https://raw.githubusercontent.com/linlang781/CLIProxyAPIPlus/main/config.example.yaml + +# 拉取并启动 +docker compose pull && docker compose up -d +``` + +### 配置说明 + +启动前请编辑 `config.yaml`： + +```yaml +# 基本配置示例 +server: + port: 8317 + +# 在此添加你的供应商配置 +``` + +### 更新到最新版本 + +```bash +cd ~/cli-proxy +docker compose pull && docker compose up -d +``` + ## 贡献 该项目仅接受第三方供应商支持的 Pull Request。任何非第三方供应商支持的 Pull Request 都将被拒绝。 From f87fe0a0e83c69c2fd58593e0babcd1a9ae06131 Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Mon, 19 Jan 2026 20:09:38 +0800 Subject: [PATCH 03/10] feat: proactive token refresh 10 minutes before expiry Amp-Thread-ID: https://ampcode.com/threads/T-019bd618-7e42-715a-960d-dd45425851e3 Co-authored-by: Amp --- sdk/auth/filestore.go | 11 ++++++++++- sdk/auth/kiro.go | 20 ++++++++++---------- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/sdk/auth/filestore.go b/sdk/auth/filestore.go index db9f7148..76361507 100644 --- a/sdk/auth/filestore.go +++ b/sdk/auth/filestore.go @@ -217,6 +217,15 @@ func (s *FileTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Auth, return nil, fmt.Errorf("stat file: %w", err) } id := s.idFor(path, baseDir) + + // Calculate NextRefreshAfter from expires_at (10 minutes before expiry) + var nextRefreshAfter time.Time + if expiresAtStr, ok := metadata["expires_at"].(string); ok && expiresAtStr != "" { + if expiresAt, err := time.Parse(time.RFC3339, expiresAtStr); err == nil { + nextRefreshAfter = expiresAt.Add(-10 * time.Minute) + } + } + auth := &cliproxyauth.Auth{ ID: id, Provider: provider, @@ -228,7 +237,7 @@ func (s *FileTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Auth, CreatedAt: info.ModTime(), UpdatedAt: info.ModTime(), LastRefreshedAt: time.Time{}, - NextRefreshAfter: time.Time{}, + NextRefreshAfter: nextRefreshAfter, } if email, ok := metadata["email"].(string); ok && email != "" { auth.Attributes["email"] = email diff --git a/sdk/auth/kiro.go b/sdk/auth/kiro.go index 7747c777..6694a217 100644 --- a/sdk/auth/kiro.go +++ b/sdk/auth/kiro.go @@ -52,9 +52,9 @@ func (a *KiroAuthenticator) Provider() string { } // RefreshLead indicates how soon before expiry a refresh should be attempted. -// Set to 5 minutes to match Antigravity and avoid frequent refresh checks while still ensuring timely token refresh. +// Set to 10 minutes for proactive refresh before token expiry. func (a *KiroAuthenticator) RefreshLead() *time.Duration { - d := 5 * time.Minute + d := 10 * time.Minute return &d } @@ -126,8 +126,8 @@ func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, UpdatedAt: now, Metadata: metadata, Attributes: attributes, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + // NextRefreshAfter: 10 minutes before expiry + NextRefreshAfter: expiresAt.Add(-10 * time.Minute), } if tokenData.Email != "" { @@ -208,8 +208,8 @@ func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.C "source": "aws-builder-id-authcode", "email": tokenData.Email, }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + // NextRefreshAfter: 10 minutes before expiry + NextRefreshAfter: expiresAt.Add(-10 * time.Minute), } if tokenData.Email != "" { @@ -292,8 +292,8 @@ func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.C "email": tokenData.Email, "region": tokenData.Region, }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), + // NextRefreshAfter: 10 minutes before expiry + NextRefreshAfter: expiresAt.Add(-10 * time.Minute), } // Display the email if extracted @@ -361,8 +361,8 @@ func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, aut updated.Metadata["refresh_token"] = tokenData.RefreshToken updated.Metadata["expires_at"] = tokenData.ExpiresAt updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization - // NextRefreshAfter is aligned with RefreshLead (5min) - updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute) + // NextRefreshAfter: 10 minutes before expiry + updated.NextRefreshAfter = expiresAt.Add(-10 * time.Minute) return updated, nil } From ace7c0ccb445c267f7e27b41f9a7df0764a0163a Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Mon, 19 Jan 2026 20:28:40 +0800 Subject: [PATCH 04/10] docs: add Kiro OAuth web authentication endpoint /v0/oauth/kiro --- README.md | 17 ++++++++++++++++- README_CN.md | 17 ++++++++++++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1555e643..092a3214 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ The Plus release stays in lockstep with the mainline features. - **OAuth Web Authentication**: Browser-based OAuth login for Kiro with beautiful web UI - **Rate Limiter**: Built-in request rate limiting to prevent API abuse -- **Background Token Refresh**: Automatic token refresh in background to avoid expiration +- **Background Token Refresh**: Automatic token refresh 10 minutes before expiration - **Metrics & Monitoring**: Request metrics collection for monitoring and debugging - **Device Fingerprint**: Device fingerprint generation for enhanced security - **Cooldown Management**: Smart cooldown mechanism for API rate limits @@ -25,6 +25,21 @@ The Plus release stays in lockstep with the mainline features. - **Model Converter**: Unified model name conversion across providers - **UTF-8 Stream Processing**: Improved streaming response handling +## Kiro Authentication + +### Web-based OAuth Login + +Access the Kiro OAuth web interface at: + +``` +http://your-server:8080/v0/oauth/kiro +``` + +This provides a browser-based OAuth flow for Kiro (AWS CodeWhisperer) authentication with: +- AWS Builder ID login +- AWS Identity Center (IDC) login +- Token import from Kiro IDE + ## Quick Deployment with Docker ### One-Command Deployment diff --git a/README_CN.md b/README_CN.md index 6ac2e483..b5b4d5f9 100644 --- a/README_CN.md +++ b/README_CN.md @@ -17,7 +17,7 @@ - **OAuth Web 认证**: 基于浏览器的 Kiro OAuth 登录，提供美观的 Web UI - **请求限流器**: 内置请求限流，防止 API 滥用 -- **后台令牌刷新**: 自动后台刷新令牌，避免过期 +- **后台令牌刷新**: 过期前 10 分钟自动刷新令牌 - **监控指标**: 请求指标收集，用于监控和调试 - **设备指纹**: 设备指纹生成，增强安全性 - **冷却管理**: 智能冷却机制，应对 API 速率限制 @@ -25,6 +25,21 @@ - **模型转换器**: 跨供应商的统一模型名称转换 - **UTF-8 流处理**: 改进的流式响应处理 +## Kiro 认证 + +### 网页端 OAuth 登录 + +访问 Kiro OAuth 网页认证界面： + +``` +http://your-server:8080/v0/oauth/kiro +``` + +提供基于浏览器的 Kiro (AWS CodeWhisperer) OAuth 认证流程，支持： +- AWS Builder ID 登录 +- AWS Identity Center (IDC) 登录 +- 从 Kiro IDE 导入令牌 + ## Docker 快速部署 ### 一键部署 From 8f06f6a9edfedc52ab42062300546bbb6c40d0db Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Mon, 19 Jan 2026 20:31:33 +0800 Subject: [PATCH 05/10] chore: remove test files containing sensitive data Amp-Thread-ID: https://ampcode.com/threads/T-019bd618-7e42-715a-960d-dd45425851e3 Co-authored-by: Amp --- test_api.py | 452 ------------------------------------------ test_auth_diff.go | 273 ------------------------- test_auth_idc_go1.go | 323 ------------------------------ test_auth_js_style.go | 237 ---------------------- test_kiro_debug.go | 348 -------------------------------- test_proxy_debug.go | 367 ---------------------------------- 6 files changed, 2000 deletions(-) delete mode 100644 test_api.py delete mode 100644 test_auth_diff.go delete mode 100644 test_auth_idc_go1.go delete mode 100644 test_auth_js_style.go delete mode 100644 test_kiro_debug.go delete mode 100644 test_proxy_debug.go diff --git a/test_api.py b/test_api.py deleted file mode 100644 index 1849e2ba..00000000 --- a/test_api.py +++ /dev/null @@ -1,452 +0,0 @@ -#!/usr/bin/env python3 -""" -CLIProxyAPI 全面测试脚本 -测试模型列表、流式输出、thinking模式及复杂任务 -""" - -import requests -import json -import time -import sys -import io -from typing import Optional, List, Dict, Any - -# 修复 Windows 控制台编码问题 -sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace') -sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace') - -# 配置 -BASE_URL = "http://localhost:8317" -API_KEY = "your-api-key-1" -HEADERS = { - "Authorization": f"Bearer {API_KEY}", - "Content-Type": "application/json" -} - -# 复杂任务提示词 - 用于测试 thinking 模式 -COMPLEX_TASK_PROMPT = """请帮我分析以下复杂的编程问题，并给出详细的解决方案： - -问题：设计一个高并发的分布式任务调度系统，需要满足以下要求： -1. 支持百万级任务队列 -2. 任务可以设置优先级、延迟执行、定时执行 -3. 支持任务依赖关系（DAG调度） -4. 失败重试机制，支持指数退避 -5. 任务结果持久化和查询 -6. 水平扩展能力 -7. 监控和告警 - -请从以下几个方面详细分析： -1. 整体架构设计 -2. 核心数据结构 -3. 调度算法选择 -4. 容错机制设计 -5. 性能优化策略 -6. 技术选型建议 - -请逐步思考每个方面，给出你的推理过程。""" - -# 简单测试提示词 -SIMPLE_PROMPT = "Hello! Please respond with 'OK' if you receive this message." - -def print_separator(title: str): - print(f"\n{'='*60}") - print(f" {title}") - print(f"{'='*60}\n") - -def print_result(name: str, success: bool, detail: str = ""): - status = "✅ PASS" if success else "❌ FAIL" - print(f"{status} | {name}") - if detail: - print(f" └─ {detail[:200]}{'...' if len(detail) > 200 else ''}") - -def get_models() -> List[str]: - """获取可用模型列表""" - print_separator("获取模型列表") - try: - resp = requests.get(f"{BASE_URL}/v1/models", headers=HEADERS, timeout=30) - if resp.status_code == 200: - data = resp.json() - models = [m.get("id", m.get("name", "unknown")) for m in data.get("data", [])] - print(f"找到 {len(models)} 个模型:") - for m in models: - print(f" - {m}") - return models - else: - print(f"❌ 获取模型列表失败: HTTP {resp.status_code}") - print(f" 响应: {resp.text[:500]}") - return [] - except Exception as e: - print(f"❌ 获取模型列表异常: {e}") - return [] - -def test_model_basic(model: str) -> tuple: - """基础可用性测试，返回 (success, error_detail)""" - try: - payload = { - "model": model, - "messages": [{"role": "user", "content": SIMPLE_PROMPT}], - "max_tokens": 50, - "stream": False - } - resp = requests.post( - f"{BASE_URL}/v1/chat/completions", - headers=HEADERS, - json=payload, - timeout=60 - ) - if resp.status_code == 200: - data = resp.json() - content = data.get("choices", [{}])[0].get("message", {}).get("content", "") - return (bool(content), f"content_len={len(content)}") - else: - return (False, f"HTTP {resp.status_code}: {resp.text[:300]}") - except Exception as e: - return (False, str(e)) - -def test_streaming(model: str) -> Dict[str, Any]: - """测试流式输出""" - result = {"success": False, "chunks": 0, "content": "", "error": None} - try: - payload = { - "model": model, - "messages": [{"role": "user", "content": "Count from 1 to 5, one number per line."}], - "max_tokens": 100, - "stream": True - } - resp = requests.post( - f"{BASE_URL}/v1/chat/completions", - headers=HEADERS, - json=payload, - timeout=60, - stream=True - ) - - if resp.status_code != 200: - result["error"] = f"HTTP {resp.status_code}: {resp.text[:200]}" - return result - - content_parts = [] - for line in resp.iter_lines(): - if line: - line_str = line.decode('utf-8') - if line_str.startswith("data: "): - data_str = line_str[6:] - if data_str.strip() == "[DONE]": - break - try: - data = json.loads(data_str) - result["chunks"] += 1 - choices = data.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - if "content" in delta and delta["content"]: - content_parts.append(delta["content"]) - except json.JSONDecodeError: - pass - except Exception as e: - result["error"] = f"Parse error: {e}, data: {data_str[:200]}" - - result["content"] = "".join(content_parts) - result["success"] = result["chunks"] > 0 and len(result["content"]) > 0 - - except Exception as e: - result["error"] = str(e) - - return result - -def test_thinking_mode(model: str, complex_task: bool = False) -> Dict[str, Any]: - """测试 thinking 模式""" - result = { - "success": False, - "has_reasoning": False, - "reasoning_content": "", - "content": "", - "error": None, - "chunks": 0 - } - - prompt = COMPLEX_TASK_PROMPT if complex_task else "What is 15 * 23? Please think step by step." - - try: - # 尝试不同的 thinking 模式参数格式 - payload = { - "model": model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": 8000 if complex_task else 2000, - "stream": True - } - - # 根据模型类型添加 thinking 参数 - if "claude" in model.lower(): - payload["thinking"] = {"type": "enabled", "budget_tokens": 5000 if complex_task else 2000} - elif "gemini" in model.lower(): - payload["thinking"] = {"thinking_budget": 5000 if complex_task else 2000} - elif "gpt" in model.lower() or "codex" in model.lower() or "o1" in model.lower() or "o3" in model.lower(): - payload["reasoning_effort"] = "high" if complex_task else "medium" - else: - # 通用格式 - payload["thinking"] = {"type": "enabled", "budget_tokens": 5000 if complex_task else 2000} - - resp = requests.post( - f"{BASE_URL}/v1/chat/completions", - headers=HEADERS, - json=payload, - timeout=300 if complex_task else 120, - stream=True - ) - - if resp.status_code != 200: - result["error"] = f"HTTP {resp.status_code}: {resp.text[:500]}" - return result - - content_parts = [] - reasoning_parts = [] - - for line in resp.iter_lines(): - if line: - line_str = line.decode('utf-8') - if line_str.startswith("data: "): - data_str = line_str[6:] - if data_str.strip() == "[DONE]": - break - try: - data = json.loads(data_str) - result["chunks"] += 1 - - choices = data.get("choices", []) - if not choices: - continue - choice = choices[0] - delta = choice.get("delta", {}) - - # 检查 reasoning_content (Claude/OpenAI格式) - if "reasoning_content" in delta and delta["reasoning_content"]: - reasoning_parts.append(delta["reasoning_content"]) - result["has_reasoning"] = True - - # 检查 thinking (Gemini格式) - if "thinking" in delta and delta["thinking"]: - reasoning_parts.append(delta["thinking"]) - result["has_reasoning"] = True - - # 常规内容 - if "content" in delta and delta["content"]: - content_parts.append(delta["content"]) - - except json.JSONDecodeError as e: - pass - except Exception as e: - result["error"] = f"Parse error: {e}" - - result["reasoning_content"] = "".join(reasoning_parts) - result["content"] = "".join(content_parts) - result["success"] = result["chunks"] > 0 and (len(result["content"]) > 0 or len(result["reasoning_content"]) > 0) - - except requests.exceptions.Timeout: - result["error"] = "Request timeout" - except Exception as e: - result["error"] = str(e) - - return result - -def run_full_test(): - """运行完整测试""" - print("\n" + "="*60) - print(" CLIProxyAPI 全面测试") - print("="*60) - print(f"目标地址: {BASE_URL}") - print(f"API Key: {API_KEY[:10]}...") - - # 1. 获取模型列表 - models = get_models() - if not models: - print("\n❌ 无法获取模型列表，测试终止") - return - - # 2. 基础可用性测试 - print_separator("基础可用性测试") - available_models = [] - for model in models: - success, detail = test_model_basic(model) - print_result(f"模型: {model}", success, detail) - if success: - available_models.append(model) - - print(f"\n可用模型: {len(available_models)}/{len(models)}") - - if not available_models: - print("\n❌ 没有可用的模型，测试终止") - return - - # 3. 流式输出测试 - print_separator("流式输出测试") - streaming_results = {} - for model in available_models: - result = test_streaming(model) - streaming_results[model] = result - detail = f"chunks={result['chunks']}, content_len={len(result['content'])}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - - # 4. Thinking 模式测试 (简单任务) - print_separator("Thinking 模式测试 (简单任务)") - thinking_results = {} - for model in available_models: - result = test_thinking_mode(model, complex_task=False) - thinking_results[model] = result - detail = f"reasoning={result['has_reasoning']}, chunks={result['chunks']}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - - # 5. Thinking 模式测试 (复杂任务) - 只测试支持 thinking 的模型 - print_separator("Thinking 模式测试 (复杂任务)") - complex_thinking_results = {} - - # 选择前3个可用模型进行复杂任务测试 - test_models = available_models[:3] - print(f"测试模型 (取前3个): {test_models}\n") - - for model in test_models: - print(f"⏳ 正在测试 {model} (复杂任务，可能需要较长时间)...") - result = test_thinking_mode(model, complex_task=True) - complex_thinking_results[model] = result - - if result["success"]: - detail = f"reasoning={result['has_reasoning']}, reasoning_len={len(result['reasoning_content'])}, content_len={len(result['content'])}" - else: - detail = f"error: {result['error']}" if result["error"] else "Unknown error" - - print_result(f"模型: {model}", result["success"], detail) - - # 如果有 reasoning 内容，打印前500字符 - if result["has_reasoning"] and result["reasoning_content"]: - print(f"\n 📝 Reasoning 内容预览 (前500字符):") - print(f" {result['reasoning_content'][:500]}...") - - # 6. 总结报告 - print_separator("测试总结报告") - - print(f"📊 模型总数: {len(models)}") - print(f"✅ 可用模型: {len(available_models)}") - print(f"❌ 不可用模型: {len(models) - len(available_models)}") - - print(f"\n📊 流式输出测试:") - streaming_pass = sum(1 for r in streaming_results.values() if r["success"]) - print(f" 通过: {streaming_pass}/{len(streaming_results)}") - - print(f"\n📊 Thinking 模式测试 (简单):") - thinking_pass = sum(1 for r in thinking_results.values() if r["success"]) - thinking_with_reasoning = sum(1 for r in thinking_results.values() if r["has_reasoning"]) - print(f" 通过: {thinking_pass}/{len(thinking_results)}") - print(f" 包含推理内容: {thinking_with_reasoning}/{len(thinking_results)}") - - print(f"\n📊 Thinking 模式测试 (复杂):") - complex_pass = sum(1 for r in complex_thinking_results.values() if r["success"]) - complex_with_reasoning = sum(1 for r in complex_thinking_results.values() if r["has_reasoning"]) - print(f" 通过: {complex_pass}/{len(complex_thinking_results)}") - print(f" 包含推理内容: {complex_with_reasoning}/{len(complex_thinking_results)}") - - # 列出所有错误 - print(f"\n📋 错误详情:") - has_errors = False - - for model, result in streaming_results.items(): - if result["error"]: - has_errors = True - print(f" [流式] {model}: {result['error'][:100]}") - - for model, result in thinking_results.items(): - if result["error"]: - has_errors = True - print(f" [Thinking简单] {model}: {result['error'][:100]}") - - for model, result in complex_thinking_results.items(): - if result["error"]: - has_errors = True - print(f" [Thinking复杂] {model}: {result['error'][:100]}") - - if not has_errors: - print(" 无错误") - - print("\n" + "="*60) - print(" 测试完成") - print("="*60 + "\n") - -def test_single_model_basic(model: str): - """单独测试一个模型的基础功能""" - print_separator(f"基础测试: {model}") - success, detail = test_model_basic(model) - print_result(f"模型: {model}", success, detail) - return success - -def test_single_model_streaming(model: str): - """单独测试一个模型的流式输出""" - print_separator(f"流式测试: {model}") - result = test_streaming(model) - detail = f"chunks={result['chunks']}, content_len={len(result['content'])}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - if result["content"]: - print(f"\n内容: {result['content'][:300]}") - return result - -def test_single_model_thinking(model: str, complex_task: bool = False): - """单独测试一个模型的thinking模式""" - task_type = "复杂" if complex_task else "简单" - print_separator(f"Thinking测试({task_type}): {model}") - result = test_thinking_mode(model, complex_task=complex_task) - detail = f"reasoning={result['has_reasoning']}, chunks={result['chunks']}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - if result["reasoning_content"]: - print(f"\nReasoning预览: {result['reasoning_content'][:500]}") - if result["content"]: - print(f"\n内容预览: {result['content'][:500]}") - return result - -def print_usage(): - print(""" -用法: python test_api.py [options] - -命令: - models - 获取模型列表 - basic - 测试单个模型基础功能 - stream - 测试单个模型流式输出 - thinking - 测试单个模型thinking模式(简单任务) - thinking-complex - 测试单个模型thinking模式(复杂任务) - all - 运行完整测试(原有功能) - -示例: - python test_api.py models - python test_api.py basic claude-sonnet - python test_api.py stream claude-sonnet - python test_api.py thinking claude-sonnet -""") - -if __name__ == "__main__": - import sys - - if len(sys.argv) < 2: - print_usage() - sys.exit(0) - - cmd = sys.argv[1].lower() - - if cmd == "models": - get_models() - elif cmd == "basic" and len(sys.argv) >= 3: - test_single_model_basic(sys.argv[2]) - elif cmd == "stream" and len(sys.argv) >= 3: - test_single_model_streaming(sys.argv[2]) - elif cmd == "thinking" and len(sys.argv) >= 3: - test_single_model_thinking(sys.argv[2], complex_task=False) - elif cmd == "thinking-complex" and len(sys.argv) >= 3: - test_single_model_thinking(sys.argv[2], complex_task=True) - elif cmd == "all": - run_full_test() - else: - print_usage() diff --git a/test_auth_diff.go b/test_auth_diff.go deleted file mode 100644 index b294622e..00000000 --- a/test_auth_diff.go +++ /dev/null @@ -1,273 +0,0 @@ -// 测试脚本 3：对比 CLIProxyAPIPlus 与官方格式的差异 -// 这个脚本分析 CLIProxyAPIPlus 保存的 token 与官方格式的差异 -// 运行方式: go run test_auth_diff.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" 测试脚本 3: Token 格式差异分析") - fmt.Println("=" + strings.Repeat("=", 59)) - - homeDir := os.Getenv("USERPROFILE") - - // 加载官方 IDE Token (Kiro IDE 生成) - fmt.Println("\n[1] 官方 Kiro IDE Token 格式") - fmt.Println("-" + strings.Repeat("-", 59)) - - ideTokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") - ideToken := loadAndAnalyze(ideTokenPath, "Kiro IDE") - - // 加载 CLIProxyAPIPlus 保存的 Token - fmt.Println("\n[2] CLIProxyAPIPlus 保存的 Token 格式") - fmt.Println("-" + strings.Repeat("-", 59)) - - cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") - files, _ := os.ReadDir(cliProxyDir) - - var cliProxyTokens []map[string]interface{} - for _, f := range files { - if strings.HasPrefix(f.Name(), "kiro") && strings.HasSuffix(f.Name(), ".json") { - p := filepath.Join(cliProxyDir, f.Name()) - token := loadAndAnalyze(p, f.Name()) - if token != nil { - cliProxyTokens = append(cliProxyTokens, token) - } - } - } - - // 对比分析 - fmt.Println("\n[3] 关键差异分析") - fmt.Println("-" + strings.Repeat("-", 59)) - - if ideToken == nil { - fmt.Println("❌ 无法加载 IDE Token，跳过对比") - } else if len(cliProxyTokens) == 0 { - fmt.Println("❌ 无法加载 CLIProxyAPIPlus Token，跳过对比") - } else { - // 对比最新的 CLIProxyAPIPlus token - cliToken := cliProxyTokens[0] - - fmt.Println("\n字段对比:") - fmt.Printf("%-20s | %-15s | %-15s\n", "字段", "IDE Token", "CLIProxy Token") - fmt.Println(strings.Repeat("-", 55)) - - fields := []string{ - "accessToken", "refreshToken", "clientId", "clientSecret", - "authMethod", "auth_method", "provider", "region", "expiresAt", "expires_at", - } - - for _, field := range fields { - ideVal := getFieldStatus(ideToken, field) - cliVal := getFieldStatus(cliToken, field) - - status := " " - if ideVal != cliVal { - if ideVal == "✅ 有" && cliVal == "❌ 无" { - status = "⚠️" - } else if ideVal == "❌ 无" && cliVal == "✅ 有" { - status = "📝" - } - } - - fmt.Printf("%-20s | %-15s | %-15s %s\n", field, ideVal, cliVal, status) - } - - // 关键问题检测 - fmt.Println("\n🔍 问题检测:") - - // 检查 clientId/clientSecret - if hasField(ideToken, "clientId") && !hasField(cliToken, "clientId") { - fmt.Println(" ⚠️ 问题: CLIProxyAPIPlus 缺少 clientId 字段!") - fmt.Println(" 原因: IdC 认证刷新 token 时需要 clientId") - } - - if hasField(ideToken, "clientSecret") && !hasField(cliToken, "clientSecret") { - fmt.Println(" ⚠️ 问题: CLIProxyAPIPlus 缺少 clientSecret 字段!") - fmt.Println(" 原因: IdC 认证刷新 token 时需要 clientSecret") - } - - // 检查字段名差异 - if hasField(cliToken, "auth_method") && !hasField(cliToken, "authMethod") { - fmt.Println(" 📝 注意: CLIProxy 使用 auth_method (snake_case)") - fmt.Println(" 而官方使用 authMethod (camelCase)") - } - - if hasField(cliToken, "expires_at") && !hasField(cliToken, "expiresAt") { - fmt.Println(" 📝 注意: CLIProxy 使用 expires_at (snake_case)") - fmt.Println(" 而官方使用 expiresAt (camelCase)") - } - } - - // Step 4: 测试使用完整格式的 token - fmt.Println("\n[4] 测试完整格式 Token (带 clientId/clientSecret)") - fmt.Println("-" + strings.Repeat("-", 59)) - - if ideToken != nil { - testWithFullToken(ideToken) - } - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 分析完成") - fmt.Println(strings.Repeat("=", 60)) - - // 给出建议 - fmt.Println("\n💡 修复建议:") - fmt.Println(" 1. CLIProxyAPIPlus 导入 token 时需要保留 clientId 和 clientSecret") - fmt.Println(" 2. IdC 认证刷新 token 必须使用这两个字段") - fmt.Println(" 3. 检查 CLIProxyAPIPlus 的 token 导入逻辑:") - fmt.Println(" - internal/auth/kiro/aws.go LoadKiroIDEToken()") - fmt.Println(" - sdk/auth/kiro.go ImportFromKiroIDE()") -} - -func loadAndAnalyze(path, name string) map[string]interface{} { - data, err := os.ReadFile(path) - if err != nil { - fmt.Printf("❌ 无法加载 %s: %v\n", name, err) - return nil - } - - var token map[string]interface{} - if err := json.Unmarshal(data, &token); err != nil { - fmt.Printf("❌ 无法解析 %s: %v\n", name, err) - return nil - } - - fmt.Printf("📄 %s\n", path) - fmt.Printf(" 字段数: %d\n", len(token)) - - // 列出所有字段 - fmt.Printf(" 字段列表: ") - keys := make([]string, 0, len(token)) - for k := range token { - keys = append(keys, k) - } - fmt.Printf("%v\n", keys) - - return token -} - -func getFieldStatus(token map[string]interface{}, field string) string { - if token == nil { - return "N/A" - } - if v, ok := token[field]; ok && v != nil && v != "" { - return "✅ 有" - } - return "❌ 无" -} - -func hasField(token map[string]interface{}, field string) bool { - if token == nil { - return false - } - v, ok := token[field] - return ok && v != nil && v != "" -} - -func testWithFullToken(token map[string]interface{}) { - accessToken, _ := token["accessToken"].(string) - refreshToken, _ := token["refreshToken"].(string) - clientId, _ := token["clientId"].(string) - clientSecret, _ := token["clientSecret"].(string) - region, _ := token["region"].(string) - - if region == "" { - region = "us-east-1" - } - - // 测试当前 accessToken - fmt.Println("\n测试当前 accessToken...") - if testAPICall(accessToken, region) { - fmt.Println("✅ 当前 accessToken 有效") - return - } - - fmt.Println("⚠️ 当前 accessToken 无效，尝试刷新...") - - // 检查是否有完整的刷新所需字段 - if clientId == "" || clientSecret == "" { - fmt.Println("❌ 缺少 clientId 或 clientSecret，无法刷新") - fmt.Println(" 这就是问题所在！") - return - } - - // 尝试刷新 - fmt.Println("\n使用完整字段刷新 token...") - url := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", region) - - requestBody := map[string]interface{}{ - "refreshToken": refreshToken, - "clientId": clientId, - "clientSecret": clientSecret, - "grantType": "refresh_token", - } - - body, _ := json.Marshal(requestBody) - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - if resp.StatusCode == 200 { - var refreshResp map[string]interface{} - json.Unmarshal(respBody, &refreshResp) - - newAccessToken, _ := refreshResp["accessToken"].(string) - fmt.Println("✅ Token 刷新成功!") - - // 验证新 token - if testAPICall(newAccessToken, region) { - fmt.Println("✅ 新 Token 验证成功!") - fmt.Println("\n✅ 结论: 使用完整格式 (含 clientId/clientSecret) 可以正常工作") - } - } else { - fmt.Printf("❌ 刷新失败: HTTP %d\n", resp.StatusCode) - fmt.Printf(" 响应: %s\n", string(respBody)) - } -} - -func testAPICall(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - return false - } - defer resp.Body.Close() - - return resp.StatusCode == 200 -} diff --git a/test_auth_idc_go1.go b/test_auth_idc_go1.go deleted file mode 100644 index 55fd5829..00000000 --- a/test_auth_idc_go1.go +++ /dev/null @@ -1,323 +0,0 @@ -// 测试脚本 1：模拟 kiro2api_go1 的 IdC 认证方式 -// 这个脚本完整模拟 kiro-gateway/temp/kiro2api_go1 的认证逻辑 -// 运行方式: go run test_auth_idc_go1.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "math/rand" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -// 配置常量 - 来自 kiro2api_go1/config/config.go -const ( - IdcRefreshTokenURL = "https://oidc.us-east-1.amazonaws.com/token" - CodeWhispererAPIURL = "https://codewhisperer.us-east-1.amazonaws.com" -) - -// AuthConfig - 来自 kiro2api_go1/auth/config.go -type AuthConfig struct { - AuthType string `json:"auth"` - RefreshToken string `json:"refreshToken"` - ClientID string `json:"clientId,omitempty"` - ClientSecret string `json:"clientSecret,omitempty"` -} - -// IdcRefreshRequest - 来自 kiro2api_go1/types/token.go -type IdcRefreshRequest struct { - ClientId string `json:"clientId"` - ClientSecret string `json:"clientSecret"` - GrantType string `json:"grantType"` - RefreshToken string `json:"refreshToken"` -} - -// RefreshResponse - 来自 kiro2api_go1/types/token.go -type RefreshResponse struct { - AccessToken string `json:"accessToken"` - RefreshToken string `json:"refreshToken,omitempty"` - ExpiresIn int `json:"expiresIn"` - TokenType string `json:"tokenType,omitempty"` -} - -// Fingerprint - 简化的指纹结构 -type Fingerprint struct { - OSType string - ConnectionBehavior string - AcceptLanguage string - SecFetchMode string - AcceptEncoding string -} - -func generateFingerprint() *Fingerprint { - osTypes := []string{"darwin", "windows", "linux"} - connections := []string{"keep-alive", "close"} - languages := []string{"en-US,en;q=0.9", "zh-CN,zh;q=0.9", "en-GB,en;q=0.9"} - fetchModes := []string{"cors", "navigate", "no-cors"} - - return &Fingerprint{ - OSType: osTypes[rand.Intn(len(osTypes))], - ConnectionBehavior: connections[rand.Intn(len(connections))], - AcceptLanguage: languages[rand.Intn(len(languages))], - SecFetchMode: fetchModes[rand.Intn(len(fetchModes))], - AcceptEncoding: "gzip, deflate, br", - } -} - -func main() { - rand.Seed(time.Now().UnixNano()) - - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" 测试脚本 1: kiro2api_go1 风格 IdC 认证") - fmt.Println("=" + strings.Repeat("=", 59)) - - // Step 1: 加载官方格式的 token 文件 - fmt.Println("\n[Step 1] 加载官方格式 Token 文件") - fmt.Println("-" + strings.Repeat("-", 59)) - - // 尝试从多个位置加载 - tokenPaths := []string{ - // 优先使用包含完整 clientId/clientSecret 的文件 - "E:/ai_project_2api/kiro-gateway/configs/kiro/kiro-auth-token-1768317098.json", - filepath.Join(os.Getenv("USERPROFILE"), ".aws", "sso", "cache", "kiro-auth-token.json"), - } - - var tokenData map[string]interface{} - var loadedPath string - - for _, p := range tokenPaths { - data, err := os.ReadFile(p) - if err == nil { - if err := json.Unmarshal(data, &tokenData); err == nil { - loadedPath = p - break - } - } - } - - if tokenData == nil { - fmt.Println("❌ 无法加载任何 token 文件") - return - } - - fmt.Printf("✅ 加载文件: %s\n", loadedPath) - - // 提取关键字段 - accessToken, _ := tokenData["accessToken"].(string) - refreshToken, _ := tokenData["refreshToken"].(string) - clientId, _ := tokenData["clientId"].(string) - clientSecret, _ := tokenData["clientSecret"].(string) - authMethod, _ := tokenData["authMethod"].(string) - region, _ := tokenData["region"].(string) - - if region == "" { - region = "us-east-1" - } - - fmt.Printf("\n当前 Token 信息:\n") - fmt.Printf(" AuthMethod: %s\n", authMethod) - fmt.Printf(" Region: %s\n", region) - fmt.Printf(" AccessToken: %s...\n", truncate(accessToken, 50)) - fmt.Printf(" RefreshToken: %s...\n", truncate(refreshToken, 50)) - fmt.Printf(" ClientID: %s\n", truncate(clientId, 30)) - fmt.Printf(" ClientSecret: %s...\n", truncate(clientSecret, 50)) - - // Step 2: 验证 IdC 认证所需字段 - fmt.Println("\n[Step 2] 验证 IdC 认证必需字段") - fmt.Println("-" + strings.Repeat("-", 59)) - - missingFields := []string{} - if refreshToken == "" { - missingFields = append(missingFields, "refreshToken") - } - if clientId == "" { - missingFields = append(missingFields, "clientId") - } - if clientSecret == "" { - missingFields = append(missingFields, "clientSecret") - } - - if len(missingFields) > 0 { - fmt.Printf("❌ 缺少必需字段: %v\n", missingFields) - fmt.Println(" IdC 认证需要: refreshToken, clientId, clientSecret") - return - } - fmt.Println("✅ 所有必需字段都存在") - - // Step 3: 测试直接使用 accessToken 调用 API - fmt.Println("\n[Step 3] 测试当前 AccessToken 有效性") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPICall(accessToken, region) { - fmt.Println("✅ 当前 AccessToken 有效，无需刷新") - } else { - fmt.Println("⚠️ 当前 AccessToken 无效，需要刷新") - - // Step 4: 使用 kiro2api_go1 风格刷新 token - fmt.Println("\n[Step 4] 使用 kiro2api_go1 风格刷新 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - newToken, err := refreshIdCToken(AuthConfig{ - AuthType: "IdC", - RefreshToken: refreshToken, - ClientID: clientId, - ClientSecret: clientSecret, - }, region) - - if err != nil { - fmt.Printf("❌ 刷新失败: %v\n", err) - return - } - - fmt.Println("✅ Token 刷新成功！") - fmt.Printf(" 新 AccessToken: %s...\n", truncate(newToken.AccessToken, 50)) - fmt.Printf(" ExpiresIn: %d 秒\n", newToken.ExpiresIn) - - // Step 5: 验证新 token - fmt.Println("\n[Step 5] 验证新 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPICall(newToken.AccessToken, region) { - fmt.Println("✅ 新 Token 验证成功！") - - // 保存新 token - saveNewToken(loadedPath, newToken, tokenData) - } else { - fmt.Println("❌ 新 Token 验证失败") - } - } - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 测试完成") - fmt.Println(strings.Repeat("=", 60)) -} - -// refreshIdCToken - 完全模拟 kiro2api_go1/auth/refresh.go 的 refreshIdCToken 函数 -func refreshIdCToken(authConfig AuthConfig, region string) (*RefreshResponse, error) { - refreshReq := IdcRefreshRequest{ - ClientId: authConfig.ClientID, - ClientSecret: authConfig.ClientSecret, - GrantType: "refresh_token", - RefreshToken: authConfig.RefreshToken, - } - - reqBody, err := json.Marshal(refreshReq) - if err != nil { - return nil, fmt.Errorf("序列化IdC请求失败: %v", err) - } - - url := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", region) - req, err := http.NewRequest("POST", url, bytes.NewBuffer(reqBody)) - if err != nil { - return nil, fmt.Errorf("创建IdC请求失败: %v", err) - } - - // 设置 IdC 特殊 headers（使用指纹随机化）- 完全模拟 kiro2api_go1 - fp := generateFingerprint() - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region)) - req.Header.Set("Connection", fp.ConnectionBehavior) - req.Header.Set("x-amz-user-agent", fmt.Sprintf("aws-sdk-js/3.738.0 ua/2.1 os/%s lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE", fp.OSType)) - req.Header.Set("Accept", "*/*") - req.Header.Set("Accept-Language", fp.AcceptLanguage) - req.Header.Set("sec-fetch-mode", fp.SecFetchMode) - req.Header.Set("User-Agent", "node") - req.Header.Set("Accept-Encoding", fp.AcceptEncoding) - - fmt.Println("发送刷新请求:") - fmt.Printf(" URL: %s\n", url) - fmt.Println(" Headers:") - for k, v := range req.Header { - if k == "Content-Type" || k == "Host" || k == "X-Amz-User-Agent" || k == "User-Agent" { - fmt.Printf(" %s: %s\n", k, v[0]) - } - } - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("IdC请求失败: %v", err) - } - defer resp.Body.Close() - - body, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("IdC刷新失败: 状态码 %d, 响应: %s", resp.StatusCode, string(body)) - } - - var refreshResp RefreshResponse - if err := json.Unmarshal(body, &refreshResp); err != nil { - return nil, fmt.Errorf("解析IdC响应失败: %v", err) - } - - return &refreshResp, nil -} - -func testAPICall(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" 请求错误: %v\n", err) - return false - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf(" API 响应: HTTP %d\n", resp.StatusCode) - - if resp.StatusCode == 200 { - return true - } - - fmt.Printf(" 错误详情: %s\n", truncate(string(respBody), 200)) - return false -} - -func saveNewToken(originalPath string, newToken *RefreshResponse, originalData map[string]interface{}) { - // 更新 token 数据 - originalData["accessToken"] = newToken.AccessToken - if newToken.RefreshToken != "" { - originalData["refreshToken"] = newToken.RefreshToken - } - originalData["expiresAt"] = time.Now().Add(time.Duration(newToken.ExpiresIn) * time.Second).Format(time.RFC3339) - - data, _ := json.MarshalIndent(originalData, "", " ") - - // 保存到新文件 - newPath := strings.TrimSuffix(originalPath, ".json") + "_refreshed.json" - if err := os.WriteFile(newPath, data, 0644); err != nil { - fmt.Printf("⚠️ 保存失败: %v\n", err) - } else { - fmt.Printf("✅ 新 Token 已保存到: %s\n", newPath) - } -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] -} diff --git a/test_auth_js_style.go b/test_auth_js_style.go deleted file mode 100644 index 6ded3305..00000000 --- a/test_auth_js_style.go +++ /dev/null @@ -1,237 +0,0 @@ -// 测试脚本 2：模拟 kiro2Api_js 的认证方式 -// 这个脚本完整模拟 kiro-gateway/temp/kiro2Api_js 的认证逻辑 -// 运行方式: go run test_auth_js_style.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -// 常量 - 来自 kiro2Api_js/src/kiro/auth.js -const ( - REFRESH_URL_TEMPLATE = "https://prod.{{region}}.auth.desktop.kiro.dev/refreshToken" - REFRESH_IDC_URL_TEMPLATE = "https://oidc.{{region}}.amazonaws.com/token" - AUTH_METHOD_SOCIAL = "social" - AUTH_METHOD_IDC = "IdC" -) - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" 测试脚本 2: kiro2Api_js 风格认证") - fmt.Println("=" + strings.Repeat("=", 59)) - - // Step 1: 加载 token 文件 - fmt.Println("\n[Step 1] 加载 Token 文件") - fmt.Println("-" + strings.Repeat("-", 59)) - - tokenPaths := []string{ - filepath.Join(os.Getenv("USERPROFILE"), ".aws", "sso", "cache", "kiro-auth-token.json"), - "E:/ai_project_2api/kiro-gateway/configs/kiro/kiro-auth-token-1768317098.json", - } - - var tokenData map[string]interface{} - var loadedPath string - - for _, p := range tokenPaths { - data, err := os.ReadFile(p) - if err == nil { - if err := json.Unmarshal(data, &tokenData); err == nil { - loadedPath = p - break - } - } - } - - if tokenData == nil { - fmt.Println("❌ 无法加载任何 token 文件") - return - } - - fmt.Printf("✅ 加载文件: %s\n", loadedPath) - - // 提取字段 - 模拟 kiro2Api_js/src/kiro/auth.js initializeAuth - accessToken, _ := tokenData["accessToken"].(string) - refreshToken, _ := tokenData["refreshToken"].(string) - clientId, _ := tokenData["clientId"].(string) - clientSecret, _ := tokenData["clientSecret"].(string) - authMethod, _ := tokenData["authMethod"].(string) - region, _ := tokenData["region"].(string) - - if region == "" { - region = "us-east-1" - fmt.Println("⚠️ Region 未设置，使用默认值 us-east-1") - } - - fmt.Printf("\nToken 信息:\n") - fmt.Printf(" AuthMethod: %s\n", authMethod) - fmt.Printf(" Region: %s\n", region) - fmt.Printf(" 有 ClientID: %v\n", clientId != "") - fmt.Printf(" 有 ClientSecret: %v\n", clientSecret != "") - - // Step 2: 测试当前 token - fmt.Println("\n[Step 2] 测试当前 AccessToken") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPI(accessToken, region) { - fmt.Println("✅ 当前 AccessToken 有效") - return - } - - fmt.Println("⚠️ 当前 AccessToken 无效，开始刷新...") - - // Step 3: 根据 authMethod 选择刷新方式 - 模拟 doRefreshToken - fmt.Println("\n[Step 3] 刷新 Token (JS 风格)") - fmt.Println("-" + strings.Repeat("-", 59)) - - var refreshURL string - var requestBody map[string]interface{} - - // 判断认证方式 - 模拟 kiro2Api_js auth.js doRefreshToken - if authMethod == AUTH_METHOD_SOCIAL { - // Social 认证 - refreshURL = strings.Replace(REFRESH_URL_TEMPLATE, "{{region}}", region, 1) - requestBody = map[string]interface{}{ - "refreshToken": refreshToken, - } - fmt.Println("使用 Social 认证方式") - } else { - // IdC 认证 (默认) - refreshURL = strings.Replace(REFRESH_IDC_URL_TEMPLATE, "{{region}}", region, 1) - requestBody = map[string]interface{}{ - "refreshToken": refreshToken, - "clientId": clientId, - "clientSecret": clientSecret, - "grantType": "refresh_token", - } - fmt.Println("使用 IdC 认证方式") - } - - fmt.Printf("刷新 URL: %s\n", refreshURL) - fmt.Printf("请求字段: %v\n", getKeys(requestBody)) - - // 发送刷新请求 - body, _ := json.Marshal(requestBody) - req, _ := http.NewRequest("POST", refreshURL, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - fmt.Printf("\n响应状态: HTTP %d\n", resp.StatusCode) - - if resp.StatusCode != 200 { - fmt.Printf("❌ 刷新失败: %s\n", string(respBody)) - - // 分析错误 - var errResp map[string]interface{} - if err := json.Unmarshal(respBody, &errResp); err == nil { - if errType, ok := errResp["error"].(string); ok { - fmt.Printf("错误类型: %s\n", errType) - if errType == "invalid_grant" { - fmt.Println("\n💡 提示: refresh_token 可能已过期，需要重新授权") - } - } - if errDesc, ok := errResp["error_description"].(string); ok { - fmt.Printf("错误描述: %s\n", errDesc) - } - } - return - } - - // 解析响应 - var refreshResp map[string]interface{} - json.Unmarshal(respBody, &refreshResp) - - newAccessToken, _ := refreshResp["accessToken"].(string) - newRefreshToken, _ := refreshResp["refreshToken"].(string) - expiresIn, _ := refreshResp["expiresIn"].(float64) - - fmt.Println("✅ Token 刷新成功!") - fmt.Printf(" 新 AccessToken: %s...\n", truncate(newAccessToken, 50)) - fmt.Printf(" ExpiresIn: %.0f 秒\n", expiresIn) - if newRefreshToken != "" { - fmt.Printf(" 新 RefreshToken: %s...\n", truncate(newRefreshToken, 50)) - } - - // Step 4: 验证新 token - fmt.Println("\n[Step 4] 验证新 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPI(newAccessToken, region) { - fmt.Println("✅ 新 Token 验证成功!") - - // 保存新 token - 模拟 saveCredentialsToFile - tokenData["accessToken"] = newAccessToken - if newRefreshToken != "" { - tokenData["refreshToken"] = newRefreshToken - } - tokenData["expiresAt"] = time.Now().Add(time.Duration(expiresIn) * time.Second).Format(time.RFC3339) - - saveData, _ := json.MarshalIndent(tokenData, "", " ") - newPath := strings.TrimSuffix(loadedPath, ".json") + "_js_refreshed.json" - os.WriteFile(newPath, saveData, 0644) - fmt.Printf("✅ 已保存到: %s\n", newPath) - } else { - fmt.Println("❌ 新 Token 验证失败") - } - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 测试完成") - fmt.Println(strings.Repeat("=", 60)) -} - -func testAPI(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - return false - } - defer resp.Body.Close() - - return resp.StatusCode == 200 -} - -func getKeys(m map[string]interface{}) []string { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - return keys -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] -} diff --git a/test_kiro_debug.go b/test_kiro_debug.go deleted file mode 100644 index 0fbbed6c..00000000 --- a/test_kiro_debug.go +++ /dev/null @@ -1,348 +0,0 @@ -// 独立测试脚本：排查 Kiro Token 403 错误 -// 运行方式: go run test_kiro_debug.go -package main - -import ( - "bytes" - "encoding/base64" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -// Token 结构 - 匹配 Kiro IDE 格式 -type KiroIDEToken struct { - AccessToken string `json:"accessToken"` - RefreshToken string `json:"refreshToken"` - ExpiresAt string `json:"expiresAt"` - ClientIDHash string `json:"clientIdHash,omitempty"` - AuthMethod string `json:"authMethod"` - Provider string `json:"provider"` - Region string `json:"region,omitempty"` -} - -// Token 结构 - 匹配 CLIProxyAPIPlus 格式 -type CLIProxyToken struct { - AccessToken string `json:"access_token"` - RefreshToken string `json:"refresh_token"` - ProfileArn string `json:"profile_arn"` - ExpiresAt string `json:"expires_at"` - AuthMethod string `json:"auth_method"` - Provider string `json:"provider"` - ClientID string `json:"client_id,omitempty"` - ClientSecret string `json:"client_secret,omitempty"` - Email string `json:"email,omitempty"` - Type string `json:"type"` -} - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" Kiro Token 403 错误排查工具") - fmt.Println("=" + strings.Repeat("=", 59)) - - homeDir, _ := os.UserHomeDir() - - // Step 1: 检查 Kiro IDE Token 文件 - fmt.Println("\n[Step 1] 检查 Kiro IDE Token 文件") - fmt.Println("-" + strings.Repeat("-", 59)) - - ideTokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") - ideToken, err := loadKiroIDEToken(ideTokenPath) - if err != nil { - fmt.Printf("❌ 无法加载 Kiro IDE Token: %v\n", err) - return - } - fmt.Printf("✅ Token 文件: %s\n", ideTokenPath) - fmt.Printf(" AuthMethod: %s\n", ideToken.AuthMethod) - fmt.Printf(" Provider: %s\n", ideToken.Provider) - fmt.Printf(" Region: %s\n", ideToken.Region) - fmt.Printf(" ExpiresAt: %s\n", ideToken.ExpiresAt) - fmt.Printf(" AccessToken (前50字符): %s...\n", truncate(ideToken.AccessToken, 50)) - - // Step 2: 检查 Token 过期状态 - fmt.Println("\n[Step 2] 检查 Token 过期状态") - fmt.Println("-" + strings.Repeat("-", 59)) - - expiresAt, err := parseExpiresAt(ideToken.ExpiresAt) - if err != nil { - fmt.Printf("❌ 无法解析过期时间: %v\n", err) - } else { - now := time.Now() - if now.After(expiresAt) { - fmt.Printf("❌ Token 已过期！过期时间: %s，当前时间: %s\n", expiresAt.Format(time.RFC3339), now.Format(time.RFC3339)) - } else { - remaining := expiresAt.Sub(now) - fmt.Printf("✅ Token 未过期，剩余: %s\n", remaining.Round(time.Second)) - } - } - - // Step 3: 检查 CLIProxyAPIPlus 保存的 Token - fmt.Println("\n[Step 3] 检查 CLIProxyAPIPlus 保存的 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") - files, _ := os.ReadDir(cliProxyDir) - for _, f := range files { - if strings.HasPrefix(f.Name(), "kiro") && strings.HasSuffix(f.Name(), ".json") { - filePath := filepath.Join(cliProxyDir, f.Name()) - cliToken, err := loadCLIProxyToken(filePath) - if err != nil { - fmt.Printf("❌ %s: 加载失败 - %v\n", f.Name(), err) - continue - } - fmt.Printf("📄 %s:\n", f.Name()) - fmt.Printf(" AuthMethod: %s\n", cliToken.AuthMethod) - fmt.Printf(" Provider: %s\n", cliToken.Provider) - fmt.Printf(" ExpiresAt: %s\n", cliToken.ExpiresAt) - fmt.Printf(" AccessToken (前50字符): %s...\n", truncate(cliToken.AccessToken, 50)) - - // 比较 Token - if cliToken.AccessToken == ideToken.AccessToken { - fmt.Printf(" ✅ AccessToken 与 IDE Token 一致\n") - } else { - fmt.Printf(" ⚠️ AccessToken 与 IDE Token 不一致！\n") - } - } - } - - // Step 4: 直接测试 Token 有效性 (调用 Kiro API) - fmt.Println("\n[Step 4] 直接测试 Token 有效性") - fmt.Println("-" + strings.Repeat("-", 59)) - - testTokenValidity(ideToken.AccessToken, ideToken.Region) - - // Step 5: 测试不同的请求头格式 - fmt.Println("\n[Step 5] 测试不同的请求头格式") - fmt.Println("-" + strings.Repeat("-", 59)) - - testDifferentHeaders(ideToken.AccessToken, ideToken.Region) - - // Step 6: 解析 JWT 内容 - fmt.Println("\n[Step 6] 解析 JWT Token 内容") - fmt.Println("-" + strings.Repeat("-", 59)) - - parseJWT(ideToken.AccessToken) - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 排查完成") - fmt.Println(strings.Repeat("=", 60)) -} - -func loadKiroIDEToken(path string) (*KiroIDEToken, error) { - data, err := os.ReadFile(path) - if err != nil { - return nil, err - } - var token KiroIDEToken - if err := json.Unmarshal(data, &token); err != nil { - return nil, err - } - return &token, nil -} - -func loadCLIProxyToken(path string) (*CLIProxyToken, error) { - data, err := os.ReadFile(path) - if err != nil { - return nil, err - } - var token CLIProxyToken - if err := json.Unmarshal(data, &token); err != nil { - return nil, err - } - return &token, nil -} - -func parseExpiresAt(s string) (time.Time, error) { - formats := []string{ - time.RFC3339, - "2006-01-02T15:04:05.000Z", - "2006-01-02T15:04:05Z", - } - for _, f := range formats { - if t, err := time.Parse(f, s); err == nil { - return t, nil - } - } - return time.Time{}, fmt.Errorf("无法解析时间格式: %s", s) -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] -} - -func testTokenValidity(accessToken, region string) { - if region == "" { - region = "us-east-1" - } - - // 测试 GetUsageLimits API - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - fmt.Printf("请求 URL: %s\n", url) - fmt.Printf("请求头:\n") - for k, v := range req.Header { - if k == "Authorization" { - fmt.Printf(" %s: Bearer %s...\n", k, truncate(v[0][7:], 30)) - } else { - fmt.Printf(" %s: %s\n", k, v[0]) - } - } - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf("响应状态: %d\n", resp.StatusCode) - fmt.Printf("响应内容: %s\n", string(respBody)) - - if resp.StatusCode == 200 { - fmt.Println("✅ Token 有效！") - } else if resp.StatusCode == 403 { - fmt.Println("❌ Token 无效或已过期 (403)") - } -} - -func testDifferentHeaders(accessToken, region string) { - if region == "" { - region = "us-east-1" - } - - tests := []struct { - name string - headers map[string]string - }{ - { - name: "最小请求头", - headers: map[string]string{ - "Content-Type": "application/json", - "Authorization": "Bearer " + accessToken, - }, - }, - { - name: "模拟 kiro2api_go1 风格", - headers: map[string]string{ - "Content-Type": "application/json", - "Accept": "text/event-stream", - "Authorization": "Bearer " + accessToken, - "x-amzn-kiro-agent-mode": "vibe", - "x-amzn-codewhisperer-optout": "true", - "amz-sdk-invocation-id": "test-invocation-id", - "amz-sdk-request": "attempt=1; max=3", - "x-amz-user-agent": "aws-sdk-js/1.0.27 KiroIDE-0.8.0-abc123", - "User-Agent": "aws-sdk-js/1.0.27 ua/2.1 os/windows#10.0 lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.27 m/E KiroIDE-0.8.0-abc123", - }, - }, - { - name: "模拟 CLIProxyAPIPlus 风格", - headers: map[string]string{ - "Content-Type": "application/x-amz-json-1.0", - "x-amz-target": "AmazonCodeWhispererService.GetUsageLimits", - "Authorization": "Bearer " + accessToken, - "Accept": "application/json", - "amz-sdk-invocation-id": "test-invocation-id", - "amz-sdk-request": "attempt=1; max=1", - "Connection": "close", - }, - }, - } - - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - for _, test := range tests { - fmt.Printf("\n测试: %s\n", test.name) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - for k, v := range test.headers { - req.Header.Set(k, v) - } - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" ❌ 请求失败: %v\n", err) - continue - } - - respBody, _ := io.ReadAll(resp.Body) - resp.Body.Close() - - if resp.StatusCode == 200 { - fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) - } else { - fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) - } - } -} - -func parseJWT(token string) { - parts := strings.Split(token, ".") - if len(parts) < 2 { - fmt.Println("Token 不是 JWT 格式") - return - } - - // 解码 header - headerData, err := base64.RawURLEncoding.DecodeString(parts[0]) - if err != nil { - fmt.Printf("无法解码 JWT header: %v\n", err) - } else { - var header map[string]interface{} - json.Unmarshal(headerData, &header) - fmt.Printf("JWT Header: %v\n", header) - } - - // 解码 payload - payloadData, err := base64.RawURLEncoding.DecodeString(parts[1]) - if err != nil { - fmt.Printf("无法解码 JWT payload: %v\n", err) - } else { - var payload map[string]interface{} - json.Unmarshal(payloadData, &payload) - fmt.Printf("JWT Payload:\n") - for k, v := range payload { - fmt.Printf(" %s: %v\n", k, v) - } - - // 检查过期时间 - if exp, ok := payload["exp"].(float64); ok { - expTime := time.Unix(int64(exp), 0) - if time.Now().After(expTime) { - fmt.Printf(" ⚠️ JWT 已过期! exp=%s\n", expTime.Format(time.RFC3339)) - } else { - fmt.Printf(" ✅ JWT 未过期, 剩余: %s\n", expTime.Sub(time.Now()).Round(time.Second)) - } - } - } -} diff --git a/test_proxy_debug.go b/test_proxy_debug.go deleted file mode 100644 index 82369e74..00000000 --- a/test_proxy_debug.go +++ /dev/null @@ -1,367 +0,0 @@ -// 测试脚本 2：通过 CLIProxyAPIPlus 代理层排查问题 -// 运行方式: go run test_proxy_debug.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -const ( - ProxyURL = "http://localhost:8317" - APIKey = "your-api-key-1" -) - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" CLIProxyAPIPlus 代理层问题排查") - fmt.Println("=" + strings.Repeat("=", 59)) - - // Step 1: 检查代理服务状态 - fmt.Println("\n[Step 1] 检查代理服务状态") - fmt.Println("-" + strings.Repeat("-", 59)) - - resp, err := http.Get(ProxyURL + "/health") - if err != nil { - fmt.Printf("❌ 代理服务不可达: %v\n", err) - fmt.Println("请确保服务正在运行: go run ./cmd/server/main.go") - return - } - resp.Body.Close() - fmt.Printf("✅ 代理服务正常 (HTTP %d)\n", resp.StatusCode) - - // Step 2: 获取模型列表 - fmt.Println("\n[Step 2] 获取模型列表") - fmt.Println("-" + strings.Repeat("-", 59)) - - models := getModels() - if len(models) == 0 { - fmt.Println("❌ 没有可用的模型，检查凭据加载") - checkCredentials() - return - } - fmt.Printf("✅ 找到 %d 个模型:\n", len(models)) - for _, m := range models { - fmt.Printf(" - %s\n", m) - } - - // Step 3: 测试模型请求 - 捕获详细错误 - fmt.Println("\n[Step 3] 测试模型请求（详细日志）") - fmt.Println("-" + strings.Repeat("-", 59)) - - if len(models) > 0 { - testModel := models[0] - testModelRequest(testModel) - } - - // Step 4: 检查代理内部 Token 状态 - fmt.Println("\n[Step 4] 检查代理服务加载的凭据") - fmt.Println("-" + strings.Repeat("-", 59)) - - checkProxyCredentials() - - // Step 5: 对比直接请求和代理请求 - fmt.Println("\n[Step 5] 对比直接请求 vs 代理请求") - fmt.Println("-" + strings.Repeat("-", 59)) - - compareDirectVsProxy() - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 排查完成") - fmt.Println(strings.Repeat("=", 60)) -} - -func getModels() []string { - req, _ := http.NewRequest("GET", ProxyURL+"/v1/models", nil) - req.Header.Set("Authorization", "Bearer "+APIKey) - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return nil - } - defer resp.Body.Close() - - body, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != 200 { - fmt.Printf("❌ HTTP %d: %s\n", resp.StatusCode, string(body)) - return nil - } - - var result struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - } - json.Unmarshal(body, &result) - - models := make([]string, len(result.Data)) - for i, m := range result.Data { - models[i] = m.ID - } - return models -} - -func checkCredentials() { - homeDir, _ := os.UserHomeDir() - cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") - - fmt.Printf("\n检查凭据目录: %s\n", cliProxyDir) - files, err := os.ReadDir(cliProxyDir) - if err != nil { - fmt.Printf("❌ 无法读取目录: %v\n", err) - return - } - - for _, f := range files { - if strings.HasSuffix(f.Name(), ".json") { - fmt.Printf(" 📄 %s\n", f.Name()) - } - } -} - -func testModelRequest(model string) { - fmt.Printf("测试模型: %s\n", model) - - payload := map[string]interface{}{ - "model": model, - "messages": []map[string]string{ - {"role": "user", "content": "Say 'OK' if you receive this."}, - }, - "max_tokens": 50, - "stream": false, - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", ProxyURL+"/v1/chat/completions", bytes.NewBuffer(body)) - req.Header.Set("Authorization", "Bearer "+APIKey) - req.Header.Set("Content-Type", "application/json") - - fmt.Println("\n发送请求:") - fmt.Printf(" URL: %s/v1/chat/completions\n", ProxyURL) - fmt.Printf(" Model: %s\n", model) - - client := &http.Client{Timeout: 60 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - fmt.Printf("\n响应:\n") - fmt.Printf(" Status: %d\n", resp.StatusCode) - fmt.Printf(" Headers:\n") - for k, v := range resp.Header { - fmt.Printf(" %s: %s\n", k, strings.Join(v, ", ")) - } - - // 格式化 JSON 输出 - var prettyJSON bytes.Buffer - if err := json.Indent(&prettyJSON, respBody, " ", " "); err == nil { - fmt.Printf(" Body:\n %s\n", prettyJSON.String()) - } else { - fmt.Printf(" Body: %s\n", string(respBody)) - } - - if resp.StatusCode == 200 { - fmt.Println("\n✅ 请求成功！") - } else { - fmt.Println("\n❌ 请求失败！分析错误原因...") - analyzeError(respBody) - } -} - -func analyzeError(body []byte) { - var errResp struct { - Message string `json:"message"` - Reason string `json:"reason"` - Error struct { - Message string `json:"message"` - Type string `json:"type"` - } `json:"error"` - } - json.Unmarshal(body, &errResp) - - if errResp.Message != "" { - fmt.Printf("错误消息: %s\n", errResp.Message) - } - if errResp.Reason != "" { - fmt.Printf("错误原因: %s\n", errResp.Reason) - } - if errResp.Error.Message != "" { - fmt.Printf("错误详情: %s (类型: %s)\n", errResp.Error.Message, errResp.Error.Type) - } - - // 分析常见错误 - bodyStr := string(body) - if strings.Contains(bodyStr, "bearer token") || strings.Contains(bodyStr, "invalid") { - fmt.Println("\n可能的原因:") - fmt.Println(" 1. Token 已过期 - 需要刷新") - fmt.Println(" 2. Token 格式不正确 - 检查凭据文件") - fmt.Println(" 3. 代理服务加载了旧的 Token") - } -} - -func checkProxyCredentials() { - // 尝试通过管理 API 获取凭据状态 - req, _ := http.NewRequest("GET", ProxyURL+"/v0/management/auth/list", nil) - // 使用配置中的管理密钥 admin123 - req.Header.Set("Authorization", "Bearer admin123") - - client := &http.Client{Timeout: 10 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 无法访问管理 API: %v\n", err) - return - } - defer resp.Body.Close() - - body, _ := io.ReadAll(resp.Body) - - if resp.StatusCode == 200 { - fmt.Println("管理 API 返回的凭据列表:") - var prettyJSON bytes.Buffer - if err := json.Indent(&prettyJSON, body, " ", " "); err == nil { - fmt.Printf("%s\n", prettyJSON.String()) - } else { - fmt.Printf("%s\n", string(body)) - } - } else { - fmt.Printf("管理 API 返回: HTTP %d\n", resp.StatusCode) - fmt.Printf("响应: %s\n", truncate(string(body), 200)) - } -} - -func compareDirectVsProxy() { - homeDir, _ := os.UserHomeDir() - tokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") - - data, err := os.ReadFile(tokenPath) - if err != nil { - fmt.Printf("❌ 无法读取 Token 文件: %v\n", err) - return - } - - var token struct { - AccessToken string `json:"accessToken"` - Region string `json:"region"` - } - json.Unmarshal(data, &token) - - if token.Region == "" { - token.Region = "us-east-1" - } - - // 直接请求 - fmt.Println("\n1. 直接请求 Kiro API:") - directSuccess := testDirectKiroAPI(token.AccessToken, token.Region) - - // 通过代理请求 - fmt.Println("\n2. 通过代理请求:") - proxySuccess := testProxyAPI() - - // 结论 - fmt.Println("\n结论:") - if directSuccess && !proxySuccess { - fmt.Println(" ⚠️ 直接请求成功，代理请求失败") - fmt.Println(" 问题在于 CLIProxyAPIPlus 代理层") - fmt.Println(" 可能原因:") - fmt.Println(" 1. 代理服务使用了过期的 Token") - fmt.Println(" 2. Token 刷新逻辑有问题") - fmt.Println(" 3. 请求头构造不正确") - } else if directSuccess && proxySuccess { - fmt.Println(" ✅ 两者都成功") - } else if !directSuccess && !proxySuccess { - fmt.Println(" ❌ 两者都失败 - Token 本身可能有问题") - } -} - -func testDirectKiroAPI(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" ❌ 请求失败: %v\n", err) - return false - } - defer resp.Body.Close() - - if resp.StatusCode == 200 { - fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) - return true - } - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) - return false -} - -func testProxyAPI() bool { - models := getModels() - if len(models) == 0 { - fmt.Println(" ❌ 没有可用模型") - return false - } - - payload := map[string]interface{}{ - "model": models[0], - "messages": []map[string]string{ - {"role": "user", "content": "Say OK"}, - }, - "max_tokens": 10, - "stream": false, - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", ProxyURL+"/v1/chat/completions", bytes.NewBuffer(body)) - req.Header.Set("Authorization", "Bearer "+APIKey) - req.Header.Set("Content-Type", "application/json") - - client := &http.Client{Timeout: 60 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" ❌ 请求失败: %v\n", err) - return false - } - defer resp.Body.Close() - - if resp.StatusCode == 200 { - fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) - return true - } - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) - return false -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] + "..." -} From 92fb6b012a9d863b7cedc344b8c0182c6545f28a Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Mon, 19 Jan 2026 20:55:51 +0800 Subject: [PATCH 06/10] feat(kiro): add manual token refresh button to OAuth web UI Amp-Thread-ID: https://ampcode.com/threads/T-019bd642-9806-75d8-9101-27812e0eb6ab Co-authored-by: Amp --- internal/auth/kiro/oauth_web.go | 156 ++++++++++++++++++++++ internal/auth/kiro/oauth_web_templates.go | 47 +++++++ sdk/auth/kiro.go | 16 ++- 3 files changed, 214 insertions(+), 5 deletions(-) diff --git a/internal/auth/kiro/oauth_web.go b/internal/auth/kiro/oauth_web.go index 13198516..81c24393 100644 --- a/internal/auth/kiro/oauth_web.go +++ b/internal/auth/kiro/oauth_web.go @@ -5,6 +5,7 @@ import ( "context" "crypto/rand" "encoding/base64" + "encoding/json" "fmt" "html/template" "net/http" @@ -85,6 +86,7 @@ func (h *OAuthWebHandler) RegisterRoutes(router gin.IRouter) { oauth.GET("/social/callback", h.handleSocialCallback) oauth.GET("/status", h.handleStatus) oauth.POST("/import", h.handleImportToken) + oauth.POST("/refresh", h.handleManualRefresh) } } @@ -823,3 +825,157 @@ func (h *OAuthWebHandler) handleImportToken(c *gin.Context) { "fileName": fileName, }) } + +// handleManualRefresh handles manual token refresh requests from the web UI. +// This allows users to trigger a token refresh when needed, without waiting +// for the automatic 5-second check and 10-minute-before-expiry refresh cycle. +// Uses the same refresh logic as kiro_executor.Refresh for consistency. +func (h *OAuthWebHandler) handleManualRefresh(c *gin.Context) { + authDir := "" + if h.cfg != nil && h.cfg.AuthDir != "" { + var err error + authDir, err = util.ResolveAuthDir(h.cfg.AuthDir) + if err != nil { + log.Errorf("OAuth Web: failed to resolve auth directory: %v", err) + } + } + + if authDir == "" { + home, err := os.UserHomeDir() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "success": false, + "error": "Failed to get home directory", + }) + return + } + authDir = filepath.Join(home, ".cli-proxy-api") + } + + // Find all kiro token files in the auth directory + files, err := os.ReadDir(authDir) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "success": false, + "error": fmt.Sprintf("Failed to read auth directory: %v", err), + }) + return + } + + var refreshedCount int + var errors []string + + for _, file := range files { + if file.IsDir() { + continue + } + name := file.Name() + if !strings.HasPrefix(name, "kiro-") || !strings.HasSuffix(name, ".json") { + continue + } + + filePath := filepath.Join(authDir, name) + data, err := os.ReadFile(filePath) + if err != nil { + errors = append(errors, fmt.Sprintf("%s: read error - %v", name, err)) + continue + } + + var storage KiroTokenStorage + if err := json.Unmarshal(data, &storage); err != nil { + errors = append(errors, fmt.Sprintf("%s: parse error - %v", name, err)) + continue + } + + if storage.RefreshToken == "" { + errors = append(errors, fmt.Sprintf("%s: no refresh token", name)) + continue + } + + // Refresh token using the same logic as kiro_executor.Refresh + tokenData, err := h.refreshTokenData(c.Request.Context(), &storage) + if err != nil { + errors = append(errors, fmt.Sprintf("%s: refresh failed - %v", name, err)) + continue + } + + // Update storage with new token data + storage.AccessToken = tokenData.AccessToken + if tokenData.RefreshToken != "" { + storage.RefreshToken = tokenData.RefreshToken + } + storage.ExpiresAt = tokenData.ExpiresAt + storage.LastRefresh = time.Now().Format(time.RFC3339) + if tokenData.ProfileArn != "" { + storage.ProfileArn = tokenData.ProfileArn + } + + // Write updated token back to file + updatedData, err := json.MarshalIndent(storage, "", " ") + if err != nil { + errors = append(errors, fmt.Sprintf("%s: marshal error - %v", name, err)) + continue + } + + tmpFile := filePath + ".tmp" + if err := os.WriteFile(tmpFile, updatedData, 0600); err != nil { + errors = append(errors, fmt.Sprintf("%s: write error - %v", name, err)) + continue + } + if err := os.Rename(tmpFile, filePath); err != nil { + errors = append(errors, fmt.Sprintf("%s: rename error - %v", name, err)) + continue + } + + log.Infof("OAuth Web: manually refreshed token in %s, expires at %s", name, tokenData.ExpiresAt) + refreshedCount++ + + // Notify callback if set + if h.onTokenObtained != nil { + h.onTokenObtained(tokenData) + } + } + + if refreshedCount == 0 && len(errors) > 0 { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "error": fmt.Sprintf("All refresh attempts failed: %v", errors), + }) + return + } + + response := gin.H{ + "success": true, + "message": fmt.Sprintf("Refreshed %d token(s)", refreshedCount), + "refreshedCount": refreshedCount, + } + if len(errors) > 0 { + response["warnings"] = errors + } + + c.JSON(http.StatusOK, response) +} + +// refreshTokenData refreshes a token using the appropriate method based on auth type. +// This mirrors the logic in kiro_executor.Refresh for consistency. +func (h *OAuthWebHandler) refreshTokenData(ctx context.Context, storage *KiroTokenStorage) (*KiroTokenData, error) { + ssoClient := NewSSOOIDCClient(h.cfg) + + switch { + case storage.ClientID != "" && storage.ClientSecret != "" && storage.AuthMethod == "idc" && storage.Region != "": + // IDC refresh with region-specific endpoint + log.Debugf("OAuth Web: using SSO OIDC refresh for IDC (region=%s)", storage.Region) + return ssoClient.RefreshTokenWithRegion(ctx, storage.ClientID, storage.ClientSecret, storage.RefreshToken, storage.Region, storage.StartURL) + + case storage.ClientID != "" && storage.ClientSecret != "" && storage.AuthMethod == "builder-id": + // Builder ID refresh with default endpoint + log.Debugf("OAuth Web: using SSO OIDC refresh for AWS Builder ID") + return ssoClient.RefreshToken(ctx, storage.ClientID, storage.ClientSecret, storage.RefreshToken) + + default: + // Fallback to Kiro's OAuth refresh endpoint (for social auth: Google/GitHub) + log.Debugf("OAuth Web: using Kiro OAuth refresh endpoint") + oauth := NewKiroOAuth(h.cfg) + return oauth.RefreshToken(ctx, storage.RefreshToken) + } +} diff --git a/internal/auth/kiro/oauth_web_templates.go b/internal/auth/kiro/oauth_web_templates.go index 064a1ff9..228677a5 100644 --- a/internal/auth/kiro/oauth_web_templates.go +++ b/internal/auth/kiro/oauth_web_templates.go @@ -541,6 +541,9 @@ const ( } .auth-btn.manual { background: #6c757d; } .auth-btn.manual:hover { background: #5a6268; } + .auth-btn.refresh { background: #17a2b8; } + .auth-btn.refresh:hover { background: #138496; } + .auth-btn.refresh:disabled { background: #7fb3bd; cursor: not-allowed; } .manual-form { background: #f8f9fa; padding: 20px; @@ -606,6 +609,13 @@ const ( 📋 Import RefreshToken from Kiro IDE + + + +
@@ -726,6 +736,43 @@ const ( btn.textContent = '📥 Import Token'; } } + + async function manualRefresh() { + const btn = document.getElementById('refreshBtn'); + const statusEl = document.getElementById('refreshStatus'); + + btn.disabled = true; + btn.innerHTML = '⏳ Refreshing...'; + statusEl.className = 'status-message'; + statusEl.style.display = 'none'; + + try { + const response = await fetch('/v0/oauth/kiro/refresh', { + method: 'POST', + headers: { 'Content-Type': 'application/json' } + }); + + const data = await response.json(); + + if (response.ok && data.success) { + statusEl.className = 'status-message success'; + let msg = '✅ ' + data.message; + if (data.warnings && data.warnings.length > 0) { + msg += ' (Warnings: ' + data.warnings.join('; ') + ')'; + } + statusEl.textContent = msg; + } else { + statusEl.className = 'status-message error'; + statusEl.textContent = '❌ ' + (data.error || data.message || 'Refresh failed'); + } + } catch (error) { + statusEl.className = 'status-message error'; + statusEl.textContent = '❌ Network error: ' + error.message; + } finally { + btn.disabled = false; + btn.innerHTML = '🔄 Manual Refresh All Tokens'; + } + } ` diff --git a/sdk/auth/kiro.go b/sdk/auth/kiro.go index 6694a217..b0687eba 100644 --- a/sdk/auth/kiro.go +++ b/sdk/auth/kiro.go @@ -66,13 +66,19 @@ func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, expiresAt = time.Now().Add(1 * time.Hour) } - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn, tokenData.ClientID) - - // Determine label based on auth method - label := fmt.Sprintf("kiro-%s", source) + // Determine label and identifier based on auth method + var label, idPart string if tokenData.AuthMethod == "idc" { label = "kiro-idc" + // For IDC auth, always use clientID as identifier + if tokenData.ClientID != "" { + idPart = kiroauth.SanitizeEmailForFilename(tokenData.ClientID) + } else { + idPart = fmt.Sprintf("%d", time.Now().UnixNano()%100000) + } + } else { + label = fmt.Sprintf("kiro-%s", source) + idPart = extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn, tokenData.ClientID) } now := time.Now() From fa2abd560abaf7c9b8eb5aa320b229fffbf19e86 Mon Sep 17 00:00:00 2001 From: "yuechenglong.5" Date: Tue, 20 Jan 2026 10:17:39 +0800 Subject: [PATCH 07/10] =?UTF-8?q?chore:=20cherry-pick=20=E6=96=87=E6=A1=A3?= =?UTF-8?q?=E6=9B=B4=E6=96=B0=E5=92=8C=E5=88=A0=E9=99=A4=E6=B5=8B=E8=AF=95?= =?UTF-8?q?=E6=96=87=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - docs: 添加 Kiro OAuth web 认证端点说明 (ace7c0c) - chore: 删除包含敏感数据的测试文件 (8f06f6a) - 保留本地修改: refresh_manager, token_repository 等 --- README.md | 17 +- README_CN.md | 17 +- cmd/server/main.go | 8 + internal/auth/kiro/oauth_web.go | 1 + internal/auth/kiro/refresh_manager.go | 145 +++++++ internal/auth/kiro/token.go | 6 +- internal/auth/kiro/token_repository.go | 273 +++++++++++++ internal/runtime/executor/kiro_executor.go | 7 + test_api.py | 452 --------------------- test_auth_diff.go | 273 ------------- test_auth_idc_go1.go | 323 --------------- test_auth_js_style.go | 237 ----------- test_kiro_debug.go | 348 ---------------- test_proxy_debug.go | 367 ----------------- 14 files changed, 469 insertions(+), 2005 deletions(-) create mode 100644 internal/auth/kiro/refresh_manager.go create mode 100644 internal/auth/kiro/token_repository.go delete mode 100644 test_api.py delete mode 100644 test_auth_diff.go delete mode 100644 test_auth_idc_go1.go delete mode 100644 test_auth_js_style.go delete mode 100644 test_kiro_debug.go delete mode 100644 test_proxy_debug.go diff --git a/README.md b/README.md index 1555e643..092a3214 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ The Plus release stays in lockstep with the mainline features. - **OAuth Web Authentication**: Browser-based OAuth login for Kiro with beautiful web UI - **Rate Limiter**: Built-in request rate limiting to prevent API abuse -- **Background Token Refresh**: Automatic token refresh in background to avoid expiration +- **Background Token Refresh**: Automatic token refresh 10 minutes before expiration - **Metrics & Monitoring**: Request metrics collection for monitoring and debugging - **Device Fingerprint**: Device fingerprint generation for enhanced security - **Cooldown Management**: Smart cooldown mechanism for API rate limits @@ -25,6 +25,21 @@ The Plus release stays in lockstep with the mainline features. - **Model Converter**: Unified model name conversion across providers - **UTF-8 Stream Processing**: Improved streaming response handling +## Kiro Authentication + +### Web-based OAuth Login + +Access the Kiro OAuth web interface at: + +``` +http://your-server:8080/v0/oauth/kiro +``` + +This provides a browser-based OAuth flow for Kiro (AWS CodeWhisperer) authentication with: +- AWS Builder ID login +- AWS Identity Center (IDC) login +- Token import from Kiro IDE + ## Quick Deployment with Docker ### One-Command Deployment diff --git a/README_CN.md b/README_CN.md index 6ac2e483..b5b4d5f9 100644 --- a/README_CN.md +++ b/README_CN.md @@ -17,7 +17,7 @@ - **OAuth Web 认证**: 基于浏览器的 Kiro OAuth 登录，提供美观的 Web UI - **请求限流器**: 内置请求限流，防止 API 滥用 -- **后台令牌刷新**: 自动后台刷新令牌，避免过期 +- **后台令牌刷新**: 过期前 10 分钟自动刷新令牌 - **监控指标**: 请求指标收集，用于监控和调试 - **设备指纹**: 设备指纹生成，增强安全性 - **冷却管理**: 智能冷却机制，应对 API 速率限制 @@ -25,6 +25,21 @@ - **模型转换器**: 跨供应商的统一模型名称转换 - **UTF-8 流处理**: 改进的流式响应处理 +## Kiro 认证 + +### 网页端 OAuth 登录 + +访问 Kiro OAuth 网页认证界面： + +``` +http://your-server:8080/v0/oauth/kiro +``` + +提供基于浏览器的 Kiro (AWS CodeWhisperer) OAuth 认证流程，支持： +- AWS Builder ID 登录 +- AWS Identity Center (IDC) 登录 +- 从 Kiro IDE 导入令牌 + ## Docker 快速部署 ### 一键部署 diff --git a/cmd/server/main.go b/cmd/server/main.go index 8148ceee..d0f70f67 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -17,6 +17,7 @@ import ( "github.com/joho/godotenv" configaccess "github.com/router-for-me/CLIProxyAPI/v6/internal/access/config_access" + "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro" "github.com/router-for-me/CLIProxyAPI/v6/internal/buildinfo" "github.com/router-for-me/CLIProxyAPI/v6/internal/cmd" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" @@ -533,6 +534,13 @@ func main() { } // Start the main proxy service managementasset.StartAutoUpdater(context.Background(), configFilePath) + + // 初始化并启动 Kiro token 后台刷新 + if cfg.AuthDir != "" { + kiro.InitializeAndStart(cfg.AuthDir, cfg) + defer kiro.StopGlobalRefreshManager() + } + cmd.StartService(cfg, configFilePath, password) } } diff --git a/internal/auth/kiro/oauth_web.go b/internal/auth/kiro/oauth_web.go index 13198516..4ffbb7fd 100644 --- a/internal/auth/kiro/oauth_web.go +++ b/internal/auth/kiro/oauth_web.go @@ -385,6 +385,7 @@ func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSess ClientSecret: session.clientSecret, Email: email, Region: session.region, + StartURL: session.startURL, } h.mu.Lock() diff --git a/internal/auth/kiro/refresh_manager.go b/internal/auth/kiro/refresh_manager.go new file mode 100644 index 00000000..cd27b432 --- /dev/null +++ b/internal/auth/kiro/refresh_manager.go @@ -0,0 +1,145 @@ +package kiro + +import ( + "context" + "sync" + "time" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + log "github.com/sirupsen/logrus" +) + +// RefreshManager 是后台刷新器的单例管理器 +type RefreshManager struct { + mu sync.Mutex + refresher *BackgroundRefresher + ctx context.Context + cancel context.CancelFunc + started bool +} + +var ( + globalRefreshManager *RefreshManager + managerOnce sync.Once +) + +// GetRefreshManager 获取全局刷新管理器实例 +func GetRefreshManager() *RefreshManager { + managerOnce.Do(func() { + globalRefreshManager = &RefreshManager{} + }) + return globalRefreshManager +} + +// Initialize 初始化后台刷新器 +// baseDir: token 文件所在的目录 +// cfg: 应用配置 +func (m *RefreshManager) Initialize(baseDir string, cfg *config.Config) error { + m.mu.Lock() + defer m.mu.Unlock() + + if m.started { + log.Debug("refresh manager: already initialized") + return nil + } + + if baseDir == "" { + log.Warn("refresh manager: base directory not provided, skipping initialization") + return nil + } + + // 创建 token 存储库 + repo := NewFileTokenRepository(baseDir) + + // 创建后台刷新器，配置参数 + m.refresher = NewBackgroundRefresher( + repo, + WithInterval(time.Minute), // 每分钟检查一次 + WithBatchSize(50), // 每批最多处理 50 个 token + WithConcurrency(10), // 最多 10 个并发刷新 + WithConfig(cfg), // 设置 OAuth 和 SSO 客户端 + ) + + log.Infof("refresh manager: initialized with base directory %s", baseDir) + return nil +} + +// Start 启动后台刷新 +func (m *RefreshManager) Start() { + m.mu.Lock() + defer m.mu.Unlock() + + if m.started { + log.Debug("refresh manager: already started") + return + } + + if m.refresher == nil { + log.Warn("refresh manager: not initialized, cannot start") + return + } + + m.ctx, m.cancel = context.WithCancel(context.Background()) + m.refresher.Start(m.ctx) + m.started = true + + log.Info("refresh manager: background refresh started") +} + +// Stop 停止后台刷新 +func (m *RefreshManager) Stop() { + m.mu.Lock() + defer m.mu.Unlock() + + if !m.started { + return + } + + if m.cancel != nil { + m.cancel() + } + + if m.refresher != nil { + m.refresher.Stop() + } + + m.started = false + log.Info("refresh manager: background refresh stopped") +} + +// IsRunning 检查后台刷新是否正在运行 +func (m *RefreshManager) IsRunning() bool { + m.mu.Lock() + defer m.mu.Unlock() + return m.started +} + +// UpdateBaseDir 更新 token 目录（用于运行时配置更改） +func (m *RefreshManager) UpdateBaseDir(baseDir string) { + m.mu.Lock() + defer m.mu.Unlock() + + if m.refresher != nil && m.refresher.tokenRepo != nil { + if repo, ok := m.refresher.tokenRepo.(*FileTokenRepository); ok { + repo.SetBaseDir(baseDir) + log.Infof("refresh manager: updated base directory to %s", baseDir) + } + } +} + +// InitializeAndStart 初始化并启动后台刷新（便捷方法） +func InitializeAndStart(baseDir string, cfg *config.Config) { + manager := GetRefreshManager() + if err := manager.Initialize(baseDir, cfg); err != nil { + log.Errorf("refresh manager: initialization failed: %v", err) + return + } + manager.Start() +} + +// StopGlobalRefreshManager 停止全局刷新管理器 +func StopGlobalRefreshManager() { + if globalRefreshManager != nil { + globalRefreshManager.Stop() + } +} diff --git a/internal/auth/kiro/token.go b/internal/auth/kiro/token.go index bfbdc795..0484a2dc 100644 --- a/internal/auth/kiro/token.go +++ b/internal/auth/kiro/token.go @@ -26,13 +26,13 @@ type KiroTokenStorage struct { // LastRefresh is the timestamp of the last token refresh LastRefresh string `json:"last_refresh"` // ClientID is the OAuth client ID (required for token refresh) - ClientID string `json:"clientId,omitempty"` + ClientID string `json:"client_id,omitempty"` // ClientSecret is the OAuth client secret (required for token refresh) - ClientSecret string `json:"clientSecret,omitempty"` + ClientSecret string `json:"client_secret,omitempty"` // Region is the AWS region Region string `json:"region,omitempty"` // StartURL is the AWS Identity Center start URL (for IDC auth) - StartURL string `json:"startUrl,omitempty"` + StartURL string `json:"start_url,omitempty"` // Email is the user's email address Email string `json:"email,omitempty"` } diff --git a/internal/auth/kiro/token_repository.go b/internal/auth/kiro/token_repository.go new file mode 100644 index 00000000..f7ed76a8 --- /dev/null +++ b/internal/auth/kiro/token_repository.go @@ -0,0 +1,273 @@ +package kiro + +import ( + "context" + "encoding/json" + "fmt" + "io/fs" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" + + log "github.com/sirupsen/logrus" +) + +// FileTokenRepository 实现 TokenRepository 接口，基于文件系统存储 +type FileTokenRepository struct { + mu sync.RWMutex + baseDir string +} + +// NewFileTokenRepository 创建一个新的文件 token 存储库 +func NewFileTokenRepository(baseDir string) *FileTokenRepository { + return &FileTokenRepository{ + baseDir: baseDir, + } +} + +// SetBaseDir 设置基础目录 +func (r *FileTokenRepository) SetBaseDir(dir string) { + r.mu.Lock() + r.baseDir = strings.TrimSpace(dir) + r.mu.Unlock() +} + +// FindOldestUnverified 查找需要刷新的 token（按最后验证时间排序） +func (r *FileTokenRepository) FindOldestUnverified(limit int) []*Token { + r.mu.RLock() + baseDir := r.baseDir + r.mu.RUnlock() + + if baseDir == "" { + log.Debug("token repository: base directory not configured") + return nil + } + + var tokens []*Token + + err := filepath.WalkDir(baseDir, func(path string, d fs.DirEntry, walkErr error) error { + if walkErr != nil { + return nil // 忽略错误，继续遍历 + } + if d.IsDir() { + return nil + } + if !strings.HasSuffix(strings.ToLower(d.Name()), ".json") { + return nil + } + + // 只处理 kiro 相关的 token 文件 + if !strings.HasPrefix(d.Name(), "kiro-") { + return nil + } + + token, err := r.readTokenFile(path) + if err != nil { + log.Debugf("token repository: failed to read token file %s: %v", path, err) + return nil + } + + if token != nil && token.RefreshToken != "" { + // 检查 token 是否需要刷新（过期前 5 分钟） + if token.ExpiresAt.IsZero() || time.Until(token.ExpiresAt) < 5*time.Minute { + tokens = append(tokens, token) + } + } + + return nil + }) + + if err != nil { + log.Warnf("token repository: error walking directory: %v", err) + } + + // 按最后验证时间排序（最旧的优先） + sort.Slice(tokens, func(i, j int) bool { + return tokens[i].LastVerified.Before(tokens[j].LastVerified) + }) + + // 限制返回数量 + if limit > 0 && len(tokens) > limit { + tokens = tokens[:limit] + } + + return tokens +} + +// UpdateToken 更新 token 并持久化到文件 +func (r *FileTokenRepository) UpdateToken(token *Token) error { + if token == nil { + return fmt.Errorf("token repository: token is nil") + } + + r.mu.RLock() + baseDir := r.baseDir + r.mu.RUnlock() + + if baseDir == "" { + return fmt.Errorf("token repository: base directory not configured") + } + + // 构建文件路径 + filePath := filepath.Join(baseDir, token.ID) + if !strings.HasSuffix(filePath, ".json") { + filePath += ".json" + } + + // 读取现有文件内容 + existingData := make(map[string]any) + if data, err := os.ReadFile(filePath); err == nil { + _ = json.Unmarshal(data, &existingData) + } + + // 更新字段 + existingData["access_token"] = token.AccessToken + existingData["refresh_token"] = token.RefreshToken + existingData["last_refresh"] = time.Now().Format(time.RFC3339) + + if !token.ExpiresAt.IsZero() { + existingData["expires_at"] = token.ExpiresAt.Format(time.RFC3339) + } + + // 保持原有的关键字段 + if token.ClientID != "" { + existingData["client_id"] = token.ClientID + } + if token.ClientSecret != "" { + existingData["client_secret"] = token.ClientSecret + } + if token.AuthMethod != "" { + existingData["auth_method"] = token.AuthMethod + } + if token.Region != "" { + existingData["region"] = token.Region + } + if token.StartURL != "" { + existingData["start_url"] = token.StartURL + } + + // 序列化并写入文件 + raw, err := json.MarshalIndent(existingData, "", " ") + if err != nil { + return fmt.Errorf("token repository: marshal failed: %w", err) + } + + // 原子写入：先写入临时文件，再重命名 + tmpPath := filePath + ".tmp" + if err := os.WriteFile(tmpPath, raw, 0o600); err != nil { + return fmt.Errorf("token repository: write temp file failed: %w", err) + } + if err := os.Rename(tmpPath, filePath); err != nil { + _ = os.Remove(tmpPath) + return fmt.Errorf("token repository: rename failed: %w", err) + } + + log.Debugf("token repository: updated token %s", token.ID) + return nil +} + +// readTokenFile 从文件读取 token +func (r *FileTokenRepository) readTokenFile(path string) (*Token, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + var metadata map[string]any + if err := json.Unmarshal(data, &metadata); err != nil { + return nil, err + } + + // 检查是否是 kiro token + tokenType, _ := metadata["type"].(string) + if tokenType != "kiro" { + return nil, nil + } + + // 检查 auth_method + authMethod, _ := metadata["auth_method"].(string) + if authMethod != "idc" && authMethod != "builder-id" { + return nil, nil // 只处理 IDC 和 Builder ID token + } + + token := &Token{ + ID: filepath.Base(path), + AuthMethod: authMethod, + } + + // 解析各字段 + if v, ok := metadata["access_token"].(string); ok { + token.AccessToken = v + } + if v, ok := metadata["refresh_token"].(string); ok { + token.RefreshToken = v + } + if v, ok := metadata["client_id"].(string); ok { + token.ClientID = v + } + if v, ok := metadata["client_secret"].(string); ok { + token.ClientSecret = v + } + if v, ok := metadata["region"].(string); ok { + token.Region = v + } + if v, ok := metadata["start_url"].(string); ok { + token.StartURL = v + } + if v, ok := metadata["provider"].(string); ok { + token.Provider = v + } + + // 解析时间字段 + if v, ok := metadata["expires_at"].(string); ok { + if t, err := time.Parse(time.RFC3339, v); err == nil { + token.ExpiresAt = t + } + } + if v, ok := metadata["last_refresh"].(string); ok { + if t, err := time.Parse(time.RFC3339, v); err == nil { + token.LastVerified = t + } + } + + return token, nil +} + +// ListKiroTokens 列出所有 Kiro token（用于调试） +func (r *FileTokenRepository) ListKiroTokens(ctx context.Context) ([]*Token, error) { + r.mu.RLock() + baseDir := r.baseDir + r.mu.RUnlock() + + if baseDir == "" { + return nil, fmt.Errorf("token repository: base directory not configured") + } + + var tokens []*Token + + err := filepath.WalkDir(baseDir, func(path string, d fs.DirEntry, walkErr error) error { + if walkErr != nil { + return nil + } + if d.IsDir() { + return nil + } + if !strings.HasPrefix(d.Name(), "kiro-") || !strings.HasSuffix(d.Name(), ".json") { + return nil + } + + token, err := r.readTokenFile(path) + if err != nil { + return nil + } + if token != nil { + tokens = append(tokens, token) + } + return nil + }) + + return tokens, err +} diff --git a/internal/runtime/executor/kiro_executor.go b/internal/runtime/executor/kiro_executor.go index b0c14c61..b842d5c8 100644 --- a/internal/runtime/executor/kiro_executor.go +++ b/internal/runtime/executor/kiro_executor.go @@ -3617,6 +3617,13 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c if tokenData.ClientSecret != "" { updated.Metadata["client_secret"] = tokenData.ClientSecret } + // Preserve region and start_url for IDC token refresh + if tokenData.Region != "" { + updated.Metadata["region"] = tokenData.Region + } + if tokenData.StartURL != "" { + updated.Metadata["start_url"] = tokenData.StartURL + } if updated.Attributes == nil { updated.Attributes = make(map[string]string) diff --git a/test_api.py b/test_api.py deleted file mode 100644 index 1849e2ba..00000000 --- a/test_api.py +++ /dev/null @@ -1,452 +0,0 @@ -#!/usr/bin/env python3 -""" -CLIProxyAPI 全面测试脚本 -测试模型列表、流式输出、thinking模式及复杂任务 -""" - -import requests -import json -import time -import sys -import io -from typing import Optional, List, Dict, Any - -# 修复 Windows 控制台编码问题 -sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace') -sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace') - -# 配置 -BASE_URL = "http://localhost:8317" -API_KEY = "your-api-key-1" -HEADERS = { - "Authorization": f"Bearer {API_KEY}", - "Content-Type": "application/json" -} - -# 复杂任务提示词 - 用于测试 thinking 模式 -COMPLEX_TASK_PROMPT = """请帮我分析以下复杂的编程问题，并给出详细的解决方案： - -问题：设计一个高并发的分布式任务调度系统，需要满足以下要求： -1. 支持百万级任务队列 -2. 任务可以设置优先级、延迟执行、定时执行 -3. 支持任务依赖关系（DAG调度） -4. 失败重试机制，支持指数退避 -5. 任务结果持久化和查询 -6. 水平扩展能力 -7. 监控和告警 - -请从以下几个方面详细分析： -1. 整体架构设计 -2. 核心数据结构 -3. 调度算法选择 -4. 容错机制设计 -5. 性能优化策略 -6. 技术选型建议 - -请逐步思考每个方面，给出你的推理过程。""" - -# 简单测试提示词 -SIMPLE_PROMPT = "Hello! Please respond with 'OK' if you receive this message." - -def print_separator(title: str): - print(f"\n{'='*60}") - print(f" {title}") - print(f"{'='*60}\n") - -def print_result(name: str, success: bool, detail: str = ""): - status = "✅ PASS" if success else "❌ FAIL" - print(f"{status} | {name}") - if detail: - print(f" └─ {detail[:200]}{'...' if len(detail) > 200 else ''}") - -def get_models() -> List[str]: - """获取可用模型列表""" - print_separator("获取模型列表") - try: - resp = requests.get(f"{BASE_URL}/v1/models", headers=HEADERS, timeout=30) - if resp.status_code == 200: - data = resp.json() - models = [m.get("id", m.get("name", "unknown")) for m in data.get("data", [])] - print(f"找到 {len(models)} 个模型:") - for m in models: - print(f" - {m}") - return models - else: - print(f"❌ 获取模型列表失败: HTTP {resp.status_code}") - print(f" 响应: {resp.text[:500]}") - return [] - except Exception as e: - print(f"❌ 获取模型列表异常: {e}") - return [] - -def test_model_basic(model: str) -> tuple: - """基础可用性测试，返回 (success, error_detail)""" - try: - payload = { - "model": model, - "messages": [{"role": "user", "content": SIMPLE_PROMPT}], - "max_tokens": 50, - "stream": False - } - resp = requests.post( - f"{BASE_URL}/v1/chat/completions", - headers=HEADERS, - json=payload, - timeout=60 - ) - if resp.status_code == 200: - data = resp.json() - content = data.get("choices", [{}])[0].get("message", {}).get("content", "") - return (bool(content), f"content_len={len(content)}") - else: - return (False, f"HTTP {resp.status_code}: {resp.text[:300]}") - except Exception as e: - return (False, str(e)) - -def test_streaming(model: str) -> Dict[str, Any]: - """测试流式输出""" - result = {"success": False, "chunks": 0, "content": "", "error": None} - try: - payload = { - "model": model, - "messages": [{"role": "user", "content": "Count from 1 to 5, one number per line."}], - "max_tokens": 100, - "stream": True - } - resp = requests.post( - f"{BASE_URL}/v1/chat/completions", - headers=HEADERS, - json=payload, - timeout=60, - stream=True - ) - - if resp.status_code != 200: - result["error"] = f"HTTP {resp.status_code}: {resp.text[:200]}" - return result - - content_parts = [] - for line in resp.iter_lines(): - if line: - line_str = line.decode('utf-8') - if line_str.startswith("data: "): - data_str = line_str[6:] - if data_str.strip() == "[DONE]": - break - try: - data = json.loads(data_str) - result["chunks"] += 1 - choices = data.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - if "content" in delta and delta["content"]: - content_parts.append(delta["content"]) - except json.JSONDecodeError: - pass - except Exception as e: - result["error"] = f"Parse error: {e}, data: {data_str[:200]}" - - result["content"] = "".join(content_parts) - result["success"] = result["chunks"] > 0 and len(result["content"]) > 0 - - except Exception as e: - result["error"] = str(e) - - return result - -def test_thinking_mode(model: str, complex_task: bool = False) -> Dict[str, Any]: - """测试 thinking 模式""" - result = { - "success": False, - "has_reasoning": False, - "reasoning_content": "", - "content": "", - "error": None, - "chunks": 0 - } - - prompt = COMPLEX_TASK_PROMPT if complex_task else "What is 15 * 23? Please think step by step." - - try: - # 尝试不同的 thinking 模式参数格式 - payload = { - "model": model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": 8000 if complex_task else 2000, - "stream": True - } - - # 根据模型类型添加 thinking 参数 - if "claude" in model.lower(): - payload["thinking"] = {"type": "enabled", "budget_tokens": 5000 if complex_task else 2000} - elif "gemini" in model.lower(): - payload["thinking"] = {"thinking_budget": 5000 if complex_task else 2000} - elif "gpt" in model.lower() or "codex" in model.lower() or "o1" in model.lower() or "o3" in model.lower(): - payload["reasoning_effort"] = "high" if complex_task else "medium" - else: - # 通用格式 - payload["thinking"] = {"type": "enabled", "budget_tokens": 5000 if complex_task else 2000} - - resp = requests.post( - f"{BASE_URL}/v1/chat/completions", - headers=HEADERS, - json=payload, - timeout=300 if complex_task else 120, - stream=True - ) - - if resp.status_code != 200: - result["error"] = f"HTTP {resp.status_code}: {resp.text[:500]}" - return result - - content_parts = [] - reasoning_parts = [] - - for line in resp.iter_lines(): - if line: - line_str = line.decode('utf-8') - if line_str.startswith("data: "): - data_str = line_str[6:] - if data_str.strip() == "[DONE]": - break - try: - data = json.loads(data_str) - result["chunks"] += 1 - - choices = data.get("choices", []) - if not choices: - continue - choice = choices[0] - delta = choice.get("delta", {}) - - # 检查 reasoning_content (Claude/OpenAI格式) - if "reasoning_content" in delta and delta["reasoning_content"]: - reasoning_parts.append(delta["reasoning_content"]) - result["has_reasoning"] = True - - # 检查 thinking (Gemini格式) - if "thinking" in delta and delta["thinking"]: - reasoning_parts.append(delta["thinking"]) - result["has_reasoning"] = True - - # 常规内容 - if "content" in delta and delta["content"]: - content_parts.append(delta["content"]) - - except json.JSONDecodeError as e: - pass - except Exception as e: - result["error"] = f"Parse error: {e}" - - result["reasoning_content"] = "".join(reasoning_parts) - result["content"] = "".join(content_parts) - result["success"] = result["chunks"] > 0 and (len(result["content"]) > 0 or len(result["reasoning_content"]) > 0) - - except requests.exceptions.Timeout: - result["error"] = "Request timeout" - except Exception as e: - result["error"] = str(e) - - return result - -def run_full_test(): - """运行完整测试""" - print("\n" + "="*60) - print(" CLIProxyAPI 全面测试") - print("="*60) - print(f"目标地址: {BASE_URL}") - print(f"API Key: {API_KEY[:10]}...") - - # 1. 获取模型列表 - models = get_models() - if not models: - print("\n❌ 无法获取模型列表，测试终止") - return - - # 2. 基础可用性测试 - print_separator("基础可用性测试") - available_models = [] - for model in models: - success, detail = test_model_basic(model) - print_result(f"模型: {model}", success, detail) - if success: - available_models.append(model) - - print(f"\n可用模型: {len(available_models)}/{len(models)}") - - if not available_models: - print("\n❌ 没有可用的模型，测试终止") - return - - # 3. 流式输出测试 - print_separator("流式输出测试") - streaming_results = {} - for model in available_models: - result = test_streaming(model) - streaming_results[model] = result - detail = f"chunks={result['chunks']}, content_len={len(result['content'])}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - - # 4. Thinking 模式测试 (简单任务) - print_separator("Thinking 模式测试 (简单任务)") - thinking_results = {} - for model in available_models: - result = test_thinking_mode(model, complex_task=False) - thinking_results[model] = result - detail = f"reasoning={result['has_reasoning']}, chunks={result['chunks']}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - - # 5. Thinking 模式测试 (复杂任务) - 只测试支持 thinking 的模型 - print_separator("Thinking 模式测试 (复杂任务)") - complex_thinking_results = {} - - # 选择前3个可用模型进行复杂任务测试 - test_models = available_models[:3] - print(f"测试模型 (取前3个): {test_models}\n") - - for model in test_models: - print(f"⏳ 正在测试 {model} (复杂任务，可能需要较长时间)...") - result = test_thinking_mode(model, complex_task=True) - complex_thinking_results[model] = result - - if result["success"]: - detail = f"reasoning={result['has_reasoning']}, reasoning_len={len(result['reasoning_content'])}, content_len={len(result['content'])}" - else: - detail = f"error: {result['error']}" if result["error"] else "Unknown error" - - print_result(f"模型: {model}", result["success"], detail) - - # 如果有 reasoning 内容，打印前500字符 - if result["has_reasoning"] and result["reasoning_content"]: - print(f"\n 📝 Reasoning 内容预览 (前500字符):") - print(f" {result['reasoning_content'][:500]}...") - - # 6. 总结报告 - print_separator("测试总结报告") - - print(f"📊 模型总数: {len(models)}") - print(f"✅ 可用模型: {len(available_models)}") - print(f"❌ 不可用模型: {len(models) - len(available_models)}") - - print(f"\n📊 流式输出测试:") - streaming_pass = sum(1 for r in streaming_results.values() if r["success"]) - print(f" 通过: {streaming_pass}/{len(streaming_results)}") - - print(f"\n📊 Thinking 模式测试 (简单):") - thinking_pass = sum(1 for r in thinking_results.values() if r["success"]) - thinking_with_reasoning = sum(1 for r in thinking_results.values() if r["has_reasoning"]) - print(f" 通过: {thinking_pass}/{len(thinking_results)}") - print(f" 包含推理内容: {thinking_with_reasoning}/{len(thinking_results)}") - - print(f"\n📊 Thinking 模式测试 (复杂):") - complex_pass = sum(1 for r in complex_thinking_results.values() if r["success"]) - complex_with_reasoning = sum(1 for r in complex_thinking_results.values() if r["has_reasoning"]) - print(f" 通过: {complex_pass}/{len(complex_thinking_results)}") - print(f" 包含推理内容: {complex_with_reasoning}/{len(complex_thinking_results)}") - - # 列出所有错误 - print(f"\n📋 错误详情:") - has_errors = False - - for model, result in streaming_results.items(): - if result["error"]: - has_errors = True - print(f" [流式] {model}: {result['error'][:100]}") - - for model, result in thinking_results.items(): - if result["error"]: - has_errors = True - print(f" [Thinking简单] {model}: {result['error'][:100]}") - - for model, result in complex_thinking_results.items(): - if result["error"]: - has_errors = True - print(f" [Thinking复杂] {model}: {result['error'][:100]}") - - if not has_errors: - print(" 无错误") - - print("\n" + "="*60) - print(" 测试完成") - print("="*60 + "\n") - -def test_single_model_basic(model: str): - """单独测试一个模型的基础功能""" - print_separator(f"基础测试: {model}") - success, detail = test_model_basic(model) - print_result(f"模型: {model}", success, detail) - return success - -def test_single_model_streaming(model: str): - """单独测试一个模型的流式输出""" - print_separator(f"流式测试: {model}") - result = test_streaming(model) - detail = f"chunks={result['chunks']}, content_len={len(result['content'])}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - if result["content"]: - print(f"\n内容: {result['content'][:300]}") - return result - -def test_single_model_thinking(model: str, complex_task: bool = False): - """单独测试一个模型的thinking模式""" - task_type = "复杂" if complex_task else "简单" - print_separator(f"Thinking测试({task_type}): {model}") - result = test_thinking_mode(model, complex_task=complex_task) - detail = f"reasoning={result['has_reasoning']}, chunks={result['chunks']}" - if result["error"]: - detail = f"error: {result['error']}" - print_result(f"模型: {model}", result["success"], detail) - if result["reasoning_content"]: - print(f"\nReasoning预览: {result['reasoning_content'][:500]}") - if result["content"]: - print(f"\n内容预览: {result['content'][:500]}") - return result - -def print_usage(): - print(""" -用法: python test_api.py [options] - -命令: - models - 获取模型列表 - basic - 测试单个模型基础功能 - stream - 测试单个模型流式输出 - thinking - 测试单个模型thinking模式(简单任务) - thinking-complex - 测试单个模型thinking模式(复杂任务) - all - 运行完整测试(原有功能) - -示例: - python test_api.py models - python test_api.py basic claude-sonnet - python test_api.py stream claude-sonnet - python test_api.py thinking claude-sonnet -""") - -if __name__ == "__main__": - import sys - - if len(sys.argv) < 2: - print_usage() - sys.exit(0) - - cmd = sys.argv[1].lower() - - if cmd == "models": - get_models() - elif cmd == "basic" and len(sys.argv) >= 3: - test_single_model_basic(sys.argv[2]) - elif cmd == "stream" and len(sys.argv) >= 3: - test_single_model_streaming(sys.argv[2]) - elif cmd == "thinking" and len(sys.argv) >= 3: - test_single_model_thinking(sys.argv[2], complex_task=False) - elif cmd == "thinking-complex" and len(sys.argv) >= 3: - test_single_model_thinking(sys.argv[2], complex_task=True) - elif cmd == "all": - run_full_test() - else: - print_usage() diff --git a/test_auth_diff.go b/test_auth_diff.go deleted file mode 100644 index b294622e..00000000 --- a/test_auth_diff.go +++ /dev/null @@ -1,273 +0,0 @@ -// 测试脚本 3：对比 CLIProxyAPIPlus 与官方格式的差异 -// 这个脚本分析 CLIProxyAPIPlus 保存的 token 与官方格式的差异 -// 运行方式: go run test_auth_diff.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" 测试脚本 3: Token 格式差异分析") - fmt.Println("=" + strings.Repeat("=", 59)) - - homeDir := os.Getenv("USERPROFILE") - - // 加载官方 IDE Token (Kiro IDE 生成) - fmt.Println("\n[1] 官方 Kiro IDE Token 格式") - fmt.Println("-" + strings.Repeat("-", 59)) - - ideTokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") - ideToken := loadAndAnalyze(ideTokenPath, "Kiro IDE") - - // 加载 CLIProxyAPIPlus 保存的 Token - fmt.Println("\n[2] CLIProxyAPIPlus 保存的 Token 格式") - fmt.Println("-" + strings.Repeat("-", 59)) - - cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") - files, _ := os.ReadDir(cliProxyDir) - - var cliProxyTokens []map[string]interface{} - for _, f := range files { - if strings.HasPrefix(f.Name(), "kiro") && strings.HasSuffix(f.Name(), ".json") { - p := filepath.Join(cliProxyDir, f.Name()) - token := loadAndAnalyze(p, f.Name()) - if token != nil { - cliProxyTokens = append(cliProxyTokens, token) - } - } - } - - // 对比分析 - fmt.Println("\n[3] 关键差异分析") - fmt.Println("-" + strings.Repeat("-", 59)) - - if ideToken == nil { - fmt.Println("❌ 无法加载 IDE Token，跳过对比") - } else if len(cliProxyTokens) == 0 { - fmt.Println("❌ 无法加载 CLIProxyAPIPlus Token，跳过对比") - } else { - // 对比最新的 CLIProxyAPIPlus token - cliToken := cliProxyTokens[0] - - fmt.Println("\n字段对比:") - fmt.Printf("%-20s | %-15s | %-15s\n", "字段", "IDE Token", "CLIProxy Token") - fmt.Println(strings.Repeat("-", 55)) - - fields := []string{ - "accessToken", "refreshToken", "clientId", "clientSecret", - "authMethod", "auth_method", "provider", "region", "expiresAt", "expires_at", - } - - for _, field := range fields { - ideVal := getFieldStatus(ideToken, field) - cliVal := getFieldStatus(cliToken, field) - - status := " " - if ideVal != cliVal { - if ideVal == "✅ 有" && cliVal == "❌ 无" { - status = "⚠️" - } else if ideVal == "❌ 无" && cliVal == "✅ 有" { - status = "📝" - } - } - - fmt.Printf("%-20s | %-15s | %-15s %s\n", field, ideVal, cliVal, status) - } - - // 关键问题检测 - fmt.Println("\n🔍 问题检测:") - - // 检查 clientId/clientSecret - if hasField(ideToken, "clientId") && !hasField(cliToken, "clientId") { - fmt.Println(" ⚠️ 问题: CLIProxyAPIPlus 缺少 clientId 字段!") - fmt.Println(" 原因: IdC 认证刷新 token 时需要 clientId") - } - - if hasField(ideToken, "clientSecret") && !hasField(cliToken, "clientSecret") { - fmt.Println(" ⚠️ 问题: CLIProxyAPIPlus 缺少 clientSecret 字段!") - fmt.Println(" 原因: IdC 认证刷新 token 时需要 clientSecret") - } - - // 检查字段名差异 - if hasField(cliToken, "auth_method") && !hasField(cliToken, "authMethod") { - fmt.Println(" 📝 注意: CLIProxy 使用 auth_method (snake_case)") - fmt.Println(" 而官方使用 authMethod (camelCase)") - } - - if hasField(cliToken, "expires_at") && !hasField(cliToken, "expiresAt") { - fmt.Println(" 📝 注意: CLIProxy 使用 expires_at (snake_case)") - fmt.Println(" 而官方使用 expiresAt (camelCase)") - } - } - - // Step 4: 测试使用完整格式的 token - fmt.Println("\n[4] 测试完整格式 Token (带 clientId/clientSecret)") - fmt.Println("-" + strings.Repeat("-", 59)) - - if ideToken != nil { - testWithFullToken(ideToken) - } - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 分析完成") - fmt.Println(strings.Repeat("=", 60)) - - // 给出建议 - fmt.Println("\n💡 修复建议:") - fmt.Println(" 1. CLIProxyAPIPlus 导入 token 时需要保留 clientId 和 clientSecret") - fmt.Println(" 2. IdC 认证刷新 token 必须使用这两个字段") - fmt.Println(" 3. 检查 CLIProxyAPIPlus 的 token 导入逻辑:") - fmt.Println(" - internal/auth/kiro/aws.go LoadKiroIDEToken()") - fmt.Println(" - sdk/auth/kiro.go ImportFromKiroIDE()") -} - -func loadAndAnalyze(path, name string) map[string]interface{} { - data, err := os.ReadFile(path) - if err != nil { - fmt.Printf("❌ 无法加载 %s: %v\n", name, err) - return nil - } - - var token map[string]interface{} - if err := json.Unmarshal(data, &token); err != nil { - fmt.Printf("❌ 无法解析 %s: %v\n", name, err) - return nil - } - - fmt.Printf("📄 %s\n", path) - fmt.Printf(" 字段数: %d\n", len(token)) - - // 列出所有字段 - fmt.Printf(" 字段列表: ") - keys := make([]string, 0, len(token)) - for k := range token { - keys = append(keys, k) - } - fmt.Printf("%v\n", keys) - - return token -} - -func getFieldStatus(token map[string]interface{}, field string) string { - if token == nil { - return "N/A" - } - if v, ok := token[field]; ok && v != nil && v != "" { - return "✅ 有" - } - return "❌ 无" -} - -func hasField(token map[string]interface{}, field string) bool { - if token == nil { - return false - } - v, ok := token[field] - return ok && v != nil && v != "" -} - -func testWithFullToken(token map[string]interface{}) { - accessToken, _ := token["accessToken"].(string) - refreshToken, _ := token["refreshToken"].(string) - clientId, _ := token["clientId"].(string) - clientSecret, _ := token["clientSecret"].(string) - region, _ := token["region"].(string) - - if region == "" { - region = "us-east-1" - } - - // 测试当前 accessToken - fmt.Println("\n测试当前 accessToken...") - if testAPICall(accessToken, region) { - fmt.Println("✅ 当前 accessToken 有效") - return - } - - fmt.Println("⚠️ 当前 accessToken 无效，尝试刷新...") - - // 检查是否有完整的刷新所需字段 - if clientId == "" || clientSecret == "" { - fmt.Println("❌ 缺少 clientId 或 clientSecret，无法刷新") - fmt.Println(" 这就是问题所在！") - return - } - - // 尝试刷新 - fmt.Println("\n使用完整字段刷新 token...") - url := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", region) - - requestBody := map[string]interface{}{ - "refreshToken": refreshToken, - "clientId": clientId, - "clientSecret": clientSecret, - "grantType": "refresh_token", - } - - body, _ := json.Marshal(requestBody) - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - if resp.StatusCode == 200 { - var refreshResp map[string]interface{} - json.Unmarshal(respBody, &refreshResp) - - newAccessToken, _ := refreshResp["accessToken"].(string) - fmt.Println("✅ Token 刷新成功!") - - // 验证新 token - if testAPICall(newAccessToken, region) { - fmt.Println("✅ 新 Token 验证成功!") - fmt.Println("\n✅ 结论: 使用完整格式 (含 clientId/clientSecret) 可以正常工作") - } - } else { - fmt.Printf("❌ 刷新失败: HTTP %d\n", resp.StatusCode) - fmt.Printf(" 响应: %s\n", string(respBody)) - } -} - -func testAPICall(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - return false - } - defer resp.Body.Close() - - return resp.StatusCode == 200 -} diff --git a/test_auth_idc_go1.go b/test_auth_idc_go1.go deleted file mode 100644 index 55fd5829..00000000 --- a/test_auth_idc_go1.go +++ /dev/null @@ -1,323 +0,0 @@ -// 测试脚本 1：模拟 kiro2api_go1 的 IdC 认证方式 -// 这个脚本完整模拟 kiro-gateway/temp/kiro2api_go1 的认证逻辑 -// 运行方式: go run test_auth_idc_go1.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "math/rand" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -// 配置常量 - 来自 kiro2api_go1/config/config.go -const ( - IdcRefreshTokenURL = "https://oidc.us-east-1.amazonaws.com/token" - CodeWhispererAPIURL = "https://codewhisperer.us-east-1.amazonaws.com" -) - -// AuthConfig - 来自 kiro2api_go1/auth/config.go -type AuthConfig struct { - AuthType string `json:"auth"` - RefreshToken string `json:"refreshToken"` - ClientID string `json:"clientId,omitempty"` - ClientSecret string `json:"clientSecret,omitempty"` -} - -// IdcRefreshRequest - 来自 kiro2api_go1/types/token.go -type IdcRefreshRequest struct { - ClientId string `json:"clientId"` - ClientSecret string `json:"clientSecret"` - GrantType string `json:"grantType"` - RefreshToken string `json:"refreshToken"` -} - -// RefreshResponse - 来自 kiro2api_go1/types/token.go -type RefreshResponse struct { - AccessToken string `json:"accessToken"` - RefreshToken string `json:"refreshToken,omitempty"` - ExpiresIn int `json:"expiresIn"` - TokenType string `json:"tokenType,omitempty"` -} - -// Fingerprint - 简化的指纹结构 -type Fingerprint struct { - OSType string - ConnectionBehavior string - AcceptLanguage string - SecFetchMode string - AcceptEncoding string -} - -func generateFingerprint() *Fingerprint { - osTypes := []string{"darwin", "windows", "linux"} - connections := []string{"keep-alive", "close"} - languages := []string{"en-US,en;q=0.9", "zh-CN,zh;q=0.9", "en-GB,en;q=0.9"} - fetchModes := []string{"cors", "navigate", "no-cors"} - - return &Fingerprint{ - OSType: osTypes[rand.Intn(len(osTypes))], - ConnectionBehavior: connections[rand.Intn(len(connections))], - AcceptLanguage: languages[rand.Intn(len(languages))], - SecFetchMode: fetchModes[rand.Intn(len(fetchModes))], - AcceptEncoding: "gzip, deflate, br", - } -} - -func main() { - rand.Seed(time.Now().UnixNano()) - - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" 测试脚本 1: kiro2api_go1 风格 IdC 认证") - fmt.Println("=" + strings.Repeat("=", 59)) - - // Step 1: 加载官方格式的 token 文件 - fmt.Println("\n[Step 1] 加载官方格式 Token 文件") - fmt.Println("-" + strings.Repeat("-", 59)) - - // 尝试从多个位置加载 - tokenPaths := []string{ - // 优先使用包含完整 clientId/clientSecret 的文件 - "E:/ai_project_2api/kiro-gateway/configs/kiro/kiro-auth-token-1768317098.json", - filepath.Join(os.Getenv("USERPROFILE"), ".aws", "sso", "cache", "kiro-auth-token.json"), - } - - var tokenData map[string]interface{} - var loadedPath string - - for _, p := range tokenPaths { - data, err := os.ReadFile(p) - if err == nil { - if err := json.Unmarshal(data, &tokenData); err == nil { - loadedPath = p - break - } - } - } - - if tokenData == nil { - fmt.Println("❌ 无法加载任何 token 文件") - return - } - - fmt.Printf("✅ 加载文件: %s\n", loadedPath) - - // 提取关键字段 - accessToken, _ := tokenData["accessToken"].(string) - refreshToken, _ := tokenData["refreshToken"].(string) - clientId, _ := tokenData["clientId"].(string) - clientSecret, _ := tokenData["clientSecret"].(string) - authMethod, _ := tokenData["authMethod"].(string) - region, _ := tokenData["region"].(string) - - if region == "" { - region = "us-east-1" - } - - fmt.Printf("\n当前 Token 信息:\n") - fmt.Printf(" AuthMethod: %s\n", authMethod) - fmt.Printf(" Region: %s\n", region) - fmt.Printf(" AccessToken: %s...\n", truncate(accessToken, 50)) - fmt.Printf(" RefreshToken: %s...\n", truncate(refreshToken, 50)) - fmt.Printf(" ClientID: %s\n", truncate(clientId, 30)) - fmt.Printf(" ClientSecret: %s...\n", truncate(clientSecret, 50)) - - // Step 2: 验证 IdC 认证所需字段 - fmt.Println("\n[Step 2] 验证 IdC 认证必需字段") - fmt.Println("-" + strings.Repeat("-", 59)) - - missingFields := []string{} - if refreshToken == "" { - missingFields = append(missingFields, "refreshToken") - } - if clientId == "" { - missingFields = append(missingFields, "clientId") - } - if clientSecret == "" { - missingFields = append(missingFields, "clientSecret") - } - - if len(missingFields) > 0 { - fmt.Printf("❌ 缺少必需字段: %v\n", missingFields) - fmt.Println(" IdC 认证需要: refreshToken, clientId, clientSecret") - return - } - fmt.Println("✅ 所有必需字段都存在") - - // Step 3: 测试直接使用 accessToken 调用 API - fmt.Println("\n[Step 3] 测试当前 AccessToken 有效性") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPICall(accessToken, region) { - fmt.Println("✅ 当前 AccessToken 有效，无需刷新") - } else { - fmt.Println("⚠️ 当前 AccessToken 无效，需要刷新") - - // Step 4: 使用 kiro2api_go1 风格刷新 token - fmt.Println("\n[Step 4] 使用 kiro2api_go1 风格刷新 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - newToken, err := refreshIdCToken(AuthConfig{ - AuthType: "IdC", - RefreshToken: refreshToken, - ClientID: clientId, - ClientSecret: clientSecret, - }, region) - - if err != nil { - fmt.Printf("❌ 刷新失败: %v\n", err) - return - } - - fmt.Println("✅ Token 刷新成功！") - fmt.Printf(" 新 AccessToken: %s...\n", truncate(newToken.AccessToken, 50)) - fmt.Printf(" ExpiresIn: %d 秒\n", newToken.ExpiresIn) - - // Step 5: 验证新 token - fmt.Println("\n[Step 5] 验证新 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPICall(newToken.AccessToken, region) { - fmt.Println("✅ 新 Token 验证成功！") - - // 保存新 token - saveNewToken(loadedPath, newToken, tokenData) - } else { - fmt.Println("❌ 新 Token 验证失败") - } - } - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 测试完成") - fmt.Println(strings.Repeat("=", 60)) -} - -// refreshIdCToken - 完全模拟 kiro2api_go1/auth/refresh.go 的 refreshIdCToken 函数 -func refreshIdCToken(authConfig AuthConfig, region string) (*RefreshResponse, error) { - refreshReq := IdcRefreshRequest{ - ClientId: authConfig.ClientID, - ClientSecret: authConfig.ClientSecret, - GrantType: "refresh_token", - RefreshToken: authConfig.RefreshToken, - } - - reqBody, err := json.Marshal(refreshReq) - if err != nil { - return nil, fmt.Errorf("序列化IdC请求失败: %v", err) - } - - url := fmt.Sprintf("https://oidc.%s.amazonaws.com/token", region) - req, err := http.NewRequest("POST", url, bytes.NewBuffer(reqBody)) - if err != nil { - return nil, fmt.Errorf("创建IdC请求失败: %v", err) - } - - // 设置 IdC 特殊 headers（使用指纹随机化）- 完全模拟 kiro2api_go1 - fp := generateFingerprint() - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region)) - req.Header.Set("Connection", fp.ConnectionBehavior) - req.Header.Set("x-amz-user-agent", fmt.Sprintf("aws-sdk-js/3.738.0 ua/2.1 os/%s lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE", fp.OSType)) - req.Header.Set("Accept", "*/*") - req.Header.Set("Accept-Language", fp.AcceptLanguage) - req.Header.Set("sec-fetch-mode", fp.SecFetchMode) - req.Header.Set("User-Agent", "node") - req.Header.Set("Accept-Encoding", fp.AcceptEncoding) - - fmt.Println("发送刷新请求:") - fmt.Printf(" URL: %s\n", url) - fmt.Println(" Headers:") - for k, v := range req.Header { - if k == "Content-Type" || k == "Host" || k == "X-Amz-User-Agent" || k == "User-Agent" { - fmt.Printf(" %s: %s\n", k, v[0]) - } - } - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("IdC请求失败: %v", err) - } - defer resp.Body.Close() - - body, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("IdC刷新失败: 状态码 %d, 响应: %s", resp.StatusCode, string(body)) - } - - var refreshResp RefreshResponse - if err := json.Unmarshal(body, &refreshResp); err != nil { - return nil, fmt.Errorf("解析IdC响应失败: %v", err) - } - - return &refreshResp, nil -} - -func testAPICall(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" 请求错误: %v\n", err) - return false - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf(" API 响应: HTTP %d\n", resp.StatusCode) - - if resp.StatusCode == 200 { - return true - } - - fmt.Printf(" 错误详情: %s\n", truncate(string(respBody), 200)) - return false -} - -func saveNewToken(originalPath string, newToken *RefreshResponse, originalData map[string]interface{}) { - // 更新 token 数据 - originalData["accessToken"] = newToken.AccessToken - if newToken.RefreshToken != "" { - originalData["refreshToken"] = newToken.RefreshToken - } - originalData["expiresAt"] = time.Now().Add(time.Duration(newToken.ExpiresIn) * time.Second).Format(time.RFC3339) - - data, _ := json.MarshalIndent(originalData, "", " ") - - // 保存到新文件 - newPath := strings.TrimSuffix(originalPath, ".json") + "_refreshed.json" - if err := os.WriteFile(newPath, data, 0644); err != nil { - fmt.Printf("⚠️ 保存失败: %v\n", err) - } else { - fmt.Printf("✅ 新 Token 已保存到: %s\n", newPath) - } -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] -} diff --git a/test_auth_js_style.go b/test_auth_js_style.go deleted file mode 100644 index 6ded3305..00000000 --- a/test_auth_js_style.go +++ /dev/null @@ -1,237 +0,0 @@ -// 测试脚本 2：模拟 kiro2Api_js 的认证方式 -// 这个脚本完整模拟 kiro-gateway/temp/kiro2Api_js 的认证逻辑 -// 运行方式: go run test_auth_js_style.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -// 常量 - 来自 kiro2Api_js/src/kiro/auth.js -const ( - REFRESH_URL_TEMPLATE = "https://prod.{{region}}.auth.desktop.kiro.dev/refreshToken" - REFRESH_IDC_URL_TEMPLATE = "https://oidc.{{region}}.amazonaws.com/token" - AUTH_METHOD_SOCIAL = "social" - AUTH_METHOD_IDC = "IdC" -) - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" 测试脚本 2: kiro2Api_js 风格认证") - fmt.Println("=" + strings.Repeat("=", 59)) - - // Step 1: 加载 token 文件 - fmt.Println("\n[Step 1] 加载 Token 文件") - fmt.Println("-" + strings.Repeat("-", 59)) - - tokenPaths := []string{ - filepath.Join(os.Getenv("USERPROFILE"), ".aws", "sso", "cache", "kiro-auth-token.json"), - "E:/ai_project_2api/kiro-gateway/configs/kiro/kiro-auth-token-1768317098.json", - } - - var tokenData map[string]interface{} - var loadedPath string - - for _, p := range tokenPaths { - data, err := os.ReadFile(p) - if err == nil { - if err := json.Unmarshal(data, &tokenData); err == nil { - loadedPath = p - break - } - } - } - - if tokenData == nil { - fmt.Println("❌ 无法加载任何 token 文件") - return - } - - fmt.Printf("✅ 加载文件: %s\n", loadedPath) - - // 提取字段 - 模拟 kiro2Api_js/src/kiro/auth.js initializeAuth - accessToken, _ := tokenData["accessToken"].(string) - refreshToken, _ := tokenData["refreshToken"].(string) - clientId, _ := tokenData["clientId"].(string) - clientSecret, _ := tokenData["clientSecret"].(string) - authMethod, _ := tokenData["authMethod"].(string) - region, _ := tokenData["region"].(string) - - if region == "" { - region = "us-east-1" - fmt.Println("⚠️ Region 未设置，使用默认值 us-east-1") - } - - fmt.Printf("\nToken 信息:\n") - fmt.Printf(" AuthMethod: %s\n", authMethod) - fmt.Printf(" Region: %s\n", region) - fmt.Printf(" 有 ClientID: %v\n", clientId != "") - fmt.Printf(" 有 ClientSecret: %v\n", clientSecret != "") - - // Step 2: 测试当前 token - fmt.Println("\n[Step 2] 测试当前 AccessToken") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPI(accessToken, region) { - fmt.Println("✅ 当前 AccessToken 有效") - return - } - - fmt.Println("⚠️ 当前 AccessToken 无效，开始刷新...") - - // Step 3: 根据 authMethod 选择刷新方式 - 模拟 doRefreshToken - fmt.Println("\n[Step 3] 刷新 Token (JS 风格)") - fmt.Println("-" + strings.Repeat("-", 59)) - - var refreshURL string - var requestBody map[string]interface{} - - // 判断认证方式 - 模拟 kiro2Api_js auth.js doRefreshToken - if authMethod == AUTH_METHOD_SOCIAL { - // Social 认证 - refreshURL = strings.Replace(REFRESH_URL_TEMPLATE, "{{region}}", region, 1) - requestBody = map[string]interface{}{ - "refreshToken": refreshToken, - } - fmt.Println("使用 Social 认证方式") - } else { - // IdC 认证 (默认) - refreshURL = strings.Replace(REFRESH_IDC_URL_TEMPLATE, "{{region}}", region, 1) - requestBody = map[string]interface{}{ - "refreshToken": refreshToken, - "clientId": clientId, - "clientSecret": clientSecret, - "grantType": "refresh_token", - } - fmt.Println("使用 IdC 认证方式") - } - - fmt.Printf("刷新 URL: %s\n", refreshURL) - fmt.Printf("请求字段: %v\n", getKeys(requestBody)) - - // 发送刷新请求 - body, _ := json.Marshal(requestBody) - req, _ := http.NewRequest("POST", refreshURL, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - fmt.Printf("\n响应状态: HTTP %d\n", resp.StatusCode) - - if resp.StatusCode != 200 { - fmt.Printf("❌ 刷新失败: %s\n", string(respBody)) - - // 分析错误 - var errResp map[string]interface{} - if err := json.Unmarshal(respBody, &errResp); err == nil { - if errType, ok := errResp["error"].(string); ok { - fmt.Printf("错误类型: %s\n", errType) - if errType == "invalid_grant" { - fmt.Println("\n💡 提示: refresh_token 可能已过期，需要重新授权") - } - } - if errDesc, ok := errResp["error_description"].(string); ok { - fmt.Printf("错误描述: %s\n", errDesc) - } - } - return - } - - // 解析响应 - var refreshResp map[string]interface{} - json.Unmarshal(respBody, &refreshResp) - - newAccessToken, _ := refreshResp["accessToken"].(string) - newRefreshToken, _ := refreshResp["refreshToken"].(string) - expiresIn, _ := refreshResp["expiresIn"].(float64) - - fmt.Println("✅ Token 刷新成功!") - fmt.Printf(" 新 AccessToken: %s...\n", truncate(newAccessToken, 50)) - fmt.Printf(" ExpiresIn: %.0f 秒\n", expiresIn) - if newRefreshToken != "" { - fmt.Printf(" 新 RefreshToken: %s...\n", truncate(newRefreshToken, 50)) - } - - // Step 4: 验证新 token - fmt.Println("\n[Step 4] 验证新 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - if testAPI(newAccessToken, region) { - fmt.Println("✅ 新 Token 验证成功!") - - // 保存新 token - 模拟 saveCredentialsToFile - tokenData["accessToken"] = newAccessToken - if newRefreshToken != "" { - tokenData["refreshToken"] = newRefreshToken - } - tokenData["expiresAt"] = time.Now().Add(time.Duration(expiresIn) * time.Second).Format(time.RFC3339) - - saveData, _ := json.MarshalIndent(tokenData, "", " ") - newPath := strings.TrimSuffix(loadedPath, ".json") + "_js_refreshed.json" - os.WriteFile(newPath, saveData, 0644) - fmt.Printf("✅ 已保存到: %s\n", newPath) - } else { - fmt.Println("❌ 新 Token 验证失败") - } - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 测试完成") - fmt.Println(strings.Repeat("=", 60)) -} - -func testAPI(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - return false - } - defer resp.Body.Close() - - return resp.StatusCode == 200 -} - -func getKeys(m map[string]interface{}) []string { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - return keys -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] -} diff --git a/test_kiro_debug.go b/test_kiro_debug.go deleted file mode 100644 index 0fbbed6c..00000000 --- a/test_kiro_debug.go +++ /dev/null @@ -1,348 +0,0 @@ -// 独立测试脚本：排查 Kiro Token 403 错误 -// 运行方式: go run test_kiro_debug.go -package main - -import ( - "bytes" - "encoding/base64" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -// Token 结构 - 匹配 Kiro IDE 格式 -type KiroIDEToken struct { - AccessToken string `json:"accessToken"` - RefreshToken string `json:"refreshToken"` - ExpiresAt string `json:"expiresAt"` - ClientIDHash string `json:"clientIdHash,omitempty"` - AuthMethod string `json:"authMethod"` - Provider string `json:"provider"` - Region string `json:"region,omitempty"` -} - -// Token 结构 - 匹配 CLIProxyAPIPlus 格式 -type CLIProxyToken struct { - AccessToken string `json:"access_token"` - RefreshToken string `json:"refresh_token"` - ProfileArn string `json:"profile_arn"` - ExpiresAt string `json:"expires_at"` - AuthMethod string `json:"auth_method"` - Provider string `json:"provider"` - ClientID string `json:"client_id,omitempty"` - ClientSecret string `json:"client_secret,omitempty"` - Email string `json:"email,omitempty"` - Type string `json:"type"` -} - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" Kiro Token 403 错误排查工具") - fmt.Println("=" + strings.Repeat("=", 59)) - - homeDir, _ := os.UserHomeDir() - - // Step 1: 检查 Kiro IDE Token 文件 - fmt.Println("\n[Step 1] 检查 Kiro IDE Token 文件") - fmt.Println("-" + strings.Repeat("-", 59)) - - ideTokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") - ideToken, err := loadKiroIDEToken(ideTokenPath) - if err != nil { - fmt.Printf("❌ 无法加载 Kiro IDE Token: %v\n", err) - return - } - fmt.Printf("✅ Token 文件: %s\n", ideTokenPath) - fmt.Printf(" AuthMethod: %s\n", ideToken.AuthMethod) - fmt.Printf(" Provider: %s\n", ideToken.Provider) - fmt.Printf(" Region: %s\n", ideToken.Region) - fmt.Printf(" ExpiresAt: %s\n", ideToken.ExpiresAt) - fmt.Printf(" AccessToken (前50字符): %s...\n", truncate(ideToken.AccessToken, 50)) - - // Step 2: 检查 Token 过期状态 - fmt.Println("\n[Step 2] 检查 Token 过期状态") - fmt.Println("-" + strings.Repeat("-", 59)) - - expiresAt, err := parseExpiresAt(ideToken.ExpiresAt) - if err != nil { - fmt.Printf("❌ 无法解析过期时间: %v\n", err) - } else { - now := time.Now() - if now.After(expiresAt) { - fmt.Printf("❌ Token 已过期！过期时间: %s，当前时间: %s\n", expiresAt.Format(time.RFC3339), now.Format(time.RFC3339)) - } else { - remaining := expiresAt.Sub(now) - fmt.Printf("✅ Token 未过期，剩余: %s\n", remaining.Round(time.Second)) - } - } - - // Step 3: 检查 CLIProxyAPIPlus 保存的 Token - fmt.Println("\n[Step 3] 检查 CLIProxyAPIPlus 保存的 Token") - fmt.Println("-" + strings.Repeat("-", 59)) - - cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") - files, _ := os.ReadDir(cliProxyDir) - for _, f := range files { - if strings.HasPrefix(f.Name(), "kiro") && strings.HasSuffix(f.Name(), ".json") { - filePath := filepath.Join(cliProxyDir, f.Name()) - cliToken, err := loadCLIProxyToken(filePath) - if err != nil { - fmt.Printf("❌ %s: 加载失败 - %v\n", f.Name(), err) - continue - } - fmt.Printf("📄 %s:\n", f.Name()) - fmt.Printf(" AuthMethod: %s\n", cliToken.AuthMethod) - fmt.Printf(" Provider: %s\n", cliToken.Provider) - fmt.Printf(" ExpiresAt: %s\n", cliToken.ExpiresAt) - fmt.Printf(" AccessToken (前50字符): %s...\n", truncate(cliToken.AccessToken, 50)) - - // 比较 Token - if cliToken.AccessToken == ideToken.AccessToken { - fmt.Printf(" ✅ AccessToken 与 IDE Token 一致\n") - } else { - fmt.Printf(" ⚠️ AccessToken 与 IDE Token 不一致！\n") - } - } - } - - // Step 4: 直接测试 Token 有效性 (调用 Kiro API) - fmt.Println("\n[Step 4] 直接测试 Token 有效性") - fmt.Println("-" + strings.Repeat("-", 59)) - - testTokenValidity(ideToken.AccessToken, ideToken.Region) - - // Step 5: 测试不同的请求头格式 - fmt.Println("\n[Step 5] 测试不同的请求头格式") - fmt.Println("-" + strings.Repeat("-", 59)) - - testDifferentHeaders(ideToken.AccessToken, ideToken.Region) - - // Step 6: 解析 JWT 内容 - fmt.Println("\n[Step 6] 解析 JWT Token 内容") - fmt.Println("-" + strings.Repeat("-", 59)) - - parseJWT(ideToken.AccessToken) - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 排查完成") - fmt.Println(strings.Repeat("=", 60)) -} - -func loadKiroIDEToken(path string) (*KiroIDEToken, error) { - data, err := os.ReadFile(path) - if err != nil { - return nil, err - } - var token KiroIDEToken - if err := json.Unmarshal(data, &token); err != nil { - return nil, err - } - return &token, nil -} - -func loadCLIProxyToken(path string) (*CLIProxyToken, error) { - data, err := os.ReadFile(path) - if err != nil { - return nil, err - } - var token CLIProxyToken - if err := json.Unmarshal(data, &token); err != nil { - return nil, err - } - return &token, nil -} - -func parseExpiresAt(s string) (time.Time, error) { - formats := []string{ - time.RFC3339, - "2006-01-02T15:04:05.000Z", - "2006-01-02T15:04:05Z", - } - for _, f := range formats { - if t, err := time.Parse(f, s); err == nil { - return t, nil - } - } - return time.Time{}, fmt.Errorf("无法解析时间格式: %s", s) -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] -} - -func testTokenValidity(accessToken, region string) { - if region == "" { - region = "us-east-1" - } - - // 测试 GetUsageLimits API - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - fmt.Printf("请求 URL: %s\n", url) - fmt.Printf("请求头:\n") - for k, v := range req.Header { - if k == "Authorization" { - fmt.Printf(" %s: Bearer %s...\n", k, truncate(v[0][7:], 30)) - } else { - fmt.Printf(" %s: %s\n", k, v[0]) - } - } - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf("响应状态: %d\n", resp.StatusCode) - fmt.Printf("响应内容: %s\n", string(respBody)) - - if resp.StatusCode == 200 { - fmt.Println("✅ Token 有效！") - } else if resp.StatusCode == 403 { - fmt.Println("❌ Token 无效或已过期 (403)") - } -} - -func testDifferentHeaders(accessToken, region string) { - if region == "" { - region = "us-east-1" - } - - tests := []struct { - name string - headers map[string]string - }{ - { - name: "最小请求头", - headers: map[string]string{ - "Content-Type": "application/json", - "Authorization": "Bearer " + accessToken, - }, - }, - { - name: "模拟 kiro2api_go1 风格", - headers: map[string]string{ - "Content-Type": "application/json", - "Accept": "text/event-stream", - "Authorization": "Bearer " + accessToken, - "x-amzn-kiro-agent-mode": "vibe", - "x-amzn-codewhisperer-optout": "true", - "amz-sdk-invocation-id": "test-invocation-id", - "amz-sdk-request": "attempt=1; max=3", - "x-amz-user-agent": "aws-sdk-js/1.0.27 KiroIDE-0.8.0-abc123", - "User-Agent": "aws-sdk-js/1.0.27 ua/2.1 os/windows#10.0 lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.27 m/E KiroIDE-0.8.0-abc123", - }, - }, - { - name: "模拟 CLIProxyAPIPlus 风格", - headers: map[string]string{ - "Content-Type": "application/x-amz-json-1.0", - "x-amz-target": "AmazonCodeWhispererService.GetUsageLimits", - "Authorization": "Bearer " + accessToken, - "Accept": "application/json", - "amz-sdk-invocation-id": "test-invocation-id", - "amz-sdk-request": "attempt=1; max=1", - "Connection": "close", - }, - }, - } - - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - for _, test := range tests { - fmt.Printf("\n测试: %s\n", test.name) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - for k, v := range test.headers { - req.Header.Set(k, v) - } - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" ❌ 请求失败: %v\n", err) - continue - } - - respBody, _ := io.ReadAll(resp.Body) - resp.Body.Close() - - if resp.StatusCode == 200 { - fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) - } else { - fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) - } - } -} - -func parseJWT(token string) { - parts := strings.Split(token, ".") - if len(parts) < 2 { - fmt.Println("Token 不是 JWT 格式") - return - } - - // 解码 header - headerData, err := base64.RawURLEncoding.DecodeString(parts[0]) - if err != nil { - fmt.Printf("无法解码 JWT header: %v\n", err) - } else { - var header map[string]interface{} - json.Unmarshal(headerData, &header) - fmt.Printf("JWT Header: %v\n", header) - } - - // 解码 payload - payloadData, err := base64.RawURLEncoding.DecodeString(parts[1]) - if err != nil { - fmt.Printf("无法解码 JWT payload: %v\n", err) - } else { - var payload map[string]interface{} - json.Unmarshal(payloadData, &payload) - fmt.Printf("JWT Payload:\n") - for k, v := range payload { - fmt.Printf(" %s: %v\n", k, v) - } - - // 检查过期时间 - if exp, ok := payload["exp"].(float64); ok { - expTime := time.Unix(int64(exp), 0) - if time.Now().After(expTime) { - fmt.Printf(" ⚠️ JWT 已过期! exp=%s\n", expTime.Format(time.RFC3339)) - } else { - fmt.Printf(" ✅ JWT 未过期, 剩余: %s\n", expTime.Sub(time.Now()).Round(time.Second)) - } - } - } -} diff --git a/test_proxy_debug.go b/test_proxy_debug.go deleted file mode 100644 index 82369e74..00000000 --- a/test_proxy_debug.go +++ /dev/null @@ -1,367 +0,0 @@ -// 测试脚本 2：通过 CLIProxyAPIPlus 代理层排查问题 -// 运行方式: go run test_proxy_debug.go -package main - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "time" -) - -const ( - ProxyURL = "http://localhost:8317" - APIKey = "your-api-key-1" -) - -func main() { - fmt.Println("=" + strings.Repeat("=", 59)) - fmt.Println(" CLIProxyAPIPlus 代理层问题排查") - fmt.Println("=" + strings.Repeat("=", 59)) - - // Step 1: 检查代理服务状态 - fmt.Println("\n[Step 1] 检查代理服务状态") - fmt.Println("-" + strings.Repeat("-", 59)) - - resp, err := http.Get(ProxyURL + "/health") - if err != nil { - fmt.Printf("❌ 代理服务不可达: %v\n", err) - fmt.Println("请确保服务正在运行: go run ./cmd/server/main.go") - return - } - resp.Body.Close() - fmt.Printf("✅ 代理服务正常 (HTTP %d)\n", resp.StatusCode) - - // Step 2: 获取模型列表 - fmt.Println("\n[Step 2] 获取模型列表") - fmt.Println("-" + strings.Repeat("-", 59)) - - models := getModels() - if len(models) == 0 { - fmt.Println("❌ 没有可用的模型，检查凭据加载") - checkCredentials() - return - } - fmt.Printf("✅ 找到 %d 个模型:\n", len(models)) - for _, m := range models { - fmt.Printf(" - %s\n", m) - } - - // Step 3: 测试模型请求 - 捕获详细错误 - fmt.Println("\n[Step 3] 测试模型请求（详细日志）") - fmt.Println("-" + strings.Repeat("-", 59)) - - if len(models) > 0 { - testModel := models[0] - testModelRequest(testModel) - } - - // Step 4: 检查代理内部 Token 状态 - fmt.Println("\n[Step 4] 检查代理服务加载的凭据") - fmt.Println("-" + strings.Repeat("-", 59)) - - checkProxyCredentials() - - // Step 5: 对比直接请求和代理请求 - fmt.Println("\n[Step 5] 对比直接请求 vs 代理请求") - fmt.Println("-" + strings.Repeat("-", 59)) - - compareDirectVsProxy() - - fmt.Println("\n" + strings.Repeat("=", 60)) - fmt.Println(" 排查完成") - fmt.Println(strings.Repeat("=", 60)) -} - -func getModels() []string { - req, _ := http.NewRequest("GET", ProxyURL+"/v1/models", nil) - req.Header.Set("Authorization", "Bearer "+APIKey) - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return nil - } - defer resp.Body.Close() - - body, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != 200 { - fmt.Printf("❌ HTTP %d: %s\n", resp.StatusCode, string(body)) - return nil - } - - var result struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - } - json.Unmarshal(body, &result) - - models := make([]string, len(result.Data)) - for i, m := range result.Data { - models[i] = m.ID - } - return models -} - -func checkCredentials() { - homeDir, _ := os.UserHomeDir() - cliProxyDir := filepath.Join(homeDir, ".cli-proxy-api") - - fmt.Printf("\n检查凭据目录: %s\n", cliProxyDir) - files, err := os.ReadDir(cliProxyDir) - if err != nil { - fmt.Printf("❌ 无法读取目录: %v\n", err) - return - } - - for _, f := range files { - if strings.HasSuffix(f.Name(), ".json") { - fmt.Printf(" 📄 %s\n", f.Name()) - } - } -} - -func testModelRequest(model string) { - fmt.Printf("测试模型: %s\n", model) - - payload := map[string]interface{}{ - "model": model, - "messages": []map[string]string{ - {"role": "user", "content": "Say 'OK' if you receive this."}, - }, - "max_tokens": 50, - "stream": false, - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", ProxyURL+"/v1/chat/completions", bytes.NewBuffer(body)) - req.Header.Set("Authorization", "Bearer "+APIKey) - req.Header.Set("Content-Type", "application/json") - - fmt.Println("\n发送请求:") - fmt.Printf(" URL: %s/v1/chat/completions\n", ProxyURL) - fmt.Printf(" Model: %s\n", model) - - client := &http.Client{Timeout: 60 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 请求失败: %v\n", err) - return - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - fmt.Printf("\n响应:\n") - fmt.Printf(" Status: %d\n", resp.StatusCode) - fmt.Printf(" Headers:\n") - for k, v := range resp.Header { - fmt.Printf(" %s: %s\n", k, strings.Join(v, ", ")) - } - - // 格式化 JSON 输出 - var prettyJSON bytes.Buffer - if err := json.Indent(&prettyJSON, respBody, " ", " "); err == nil { - fmt.Printf(" Body:\n %s\n", prettyJSON.String()) - } else { - fmt.Printf(" Body: %s\n", string(respBody)) - } - - if resp.StatusCode == 200 { - fmt.Println("\n✅ 请求成功！") - } else { - fmt.Println("\n❌ 请求失败！分析错误原因...") - analyzeError(respBody) - } -} - -func analyzeError(body []byte) { - var errResp struct { - Message string `json:"message"` - Reason string `json:"reason"` - Error struct { - Message string `json:"message"` - Type string `json:"type"` - } `json:"error"` - } - json.Unmarshal(body, &errResp) - - if errResp.Message != "" { - fmt.Printf("错误消息: %s\n", errResp.Message) - } - if errResp.Reason != "" { - fmt.Printf("错误原因: %s\n", errResp.Reason) - } - if errResp.Error.Message != "" { - fmt.Printf("错误详情: %s (类型: %s)\n", errResp.Error.Message, errResp.Error.Type) - } - - // 分析常见错误 - bodyStr := string(body) - if strings.Contains(bodyStr, "bearer token") || strings.Contains(bodyStr, "invalid") { - fmt.Println("\n可能的原因:") - fmt.Println(" 1. Token 已过期 - 需要刷新") - fmt.Println(" 2. Token 格式不正确 - 检查凭据文件") - fmt.Println(" 3. 代理服务加载了旧的 Token") - } -} - -func checkProxyCredentials() { - // 尝试通过管理 API 获取凭据状态 - req, _ := http.NewRequest("GET", ProxyURL+"/v0/management/auth/list", nil) - // 使用配置中的管理密钥 admin123 - req.Header.Set("Authorization", "Bearer admin123") - - client := &http.Client{Timeout: 10 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf("❌ 无法访问管理 API: %v\n", err) - return - } - defer resp.Body.Close() - - body, _ := io.ReadAll(resp.Body) - - if resp.StatusCode == 200 { - fmt.Println("管理 API 返回的凭据列表:") - var prettyJSON bytes.Buffer - if err := json.Indent(&prettyJSON, body, " ", " "); err == nil { - fmt.Printf("%s\n", prettyJSON.String()) - } else { - fmt.Printf("%s\n", string(body)) - } - } else { - fmt.Printf("管理 API 返回: HTTP %d\n", resp.StatusCode) - fmt.Printf("响应: %s\n", truncate(string(body), 200)) - } -} - -func compareDirectVsProxy() { - homeDir, _ := os.UserHomeDir() - tokenPath := filepath.Join(homeDir, ".aws", "sso", "cache", "kiro-auth-token.json") - - data, err := os.ReadFile(tokenPath) - if err != nil { - fmt.Printf("❌ 无法读取 Token 文件: %v\n", err) - return - } - - var token struct { - AccessToken string `json:"accessToken"` - Region string `json:"region"` - } - json.Unmarshal(data, &token) - - if token.Region == "" { - token.Region = "us-east-1" - } - - // 直接请求 - fmt.Println("\n1. 直接请求 Kiro API:") - directSuccess := testDirectKiroAPI(token.AccessToken, token.Region) - - // 通过代理请求 - fmt.Println("\n2. 通过代理请求:") - proxySuccess := testProxyAPI() - - // 结论 - fmt.Println("\n结论:") - if directSuccess && !proxySuccess { - fmt.Println(" ⚠️ 直接请求成功，代理请求失败") - fmt.Println(" 问题在于 CLIProxyAPIPlus 代理层") - fmt.Println(" 可能原因:") - fmt.Println(" 1. 代理服务使用了过期的 Token") - fmt.Println(" 2. Token 刷新逻辑有问题") - fmt.Println(" 3. 请求头构造不正确") - } else if directSuccess && proxySuccess { - fmt.Println(" ✅ 两者都成功") - } else if !directSuccess && !proxySuccess { - fmt.Println(" ❌ 两者都失败 - Token 本身可能有问题") - } -} - -func testDirectKiroAPI(accessToken, region string) bool { - url := fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region) - - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - "isEmailRequired": true, - "resourceType": "AGENTIC_REQUEST", - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body)) - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.GetUsageLimits") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" ❌ 请求失败: %v\n", err) - return false - } - defer resp.Body.Close() - - if resp.StatusCode == 200 { - fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) - return true - } - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) - return false -} - -func testProxyAPI() bool { - models := getModels() - if len(models) == 0 { - fmt.Println(" ❌ 没有可用模型") - return false - } - - payload := map[string]interface{}{ - "model": models[0], - "messages": []map[string]string{ - {"role": "user", "content": "Say OK"}, - }, - "max_tokens": 10, - "stream": false, - } - body, _ := json.Marshal(payload) - - req, _ := http.NewRequest("POST", ProxyURL+"/v1/chat/completions", bytes.NewBuffer(body)) - req.Header.Set("Authorization", "Bearer "+APIKey) - req.Header.Set("Content-Type", "application/json") - - client := &http.Client{Timeout: 60 * time.Second} - resp, err := client.Do(req) - if err != nil { - fmt.Printf(" ❌ 请求失败: %v\n", err) - return false - } - defer resp.Body.Close() - - if resp.StatusCode == 200 { - fmt.Printf(" ✅ 成功 (HTTP %d)\n", resp.StatusCode) - return true - } - respBody, _ := io.ReadAll(resp.Body) - fmt.Printf(" ❌ 失败 (HTTP %d): %s\n", resp.StatusCode, truncate(string(respBody), 100)) - return false -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] + "..." -} From a9ee971e1c30499c409319471b7dea8b23f0c2ee Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Tue, 20 Jan 2026 21:57:45 +0800 Subject: [PATCH 08/10] fix(kiro): improve auto-refresh and IDC auth file handling Amp-Thread-ID: https://ampcode.com/threads/T-019bdb94-80e3-7302-be0f-a69937826d13 Co-authored-by: Amp --- internal/auth/kiro/aws_auth.go | 20 +++++++++++++++++ internal/auth/kiro/oauth_web.go | 25 +++++++++++----------- internal/runtime/executor/kiro_executor.go | 14 ++++++------ sdk/auth/filestore.go | 4 ++-- sdk/auth/kiro.go | 20 ++++++++--------- sdk/cliproxy/auth/conductor.go | 2 +- 6 files changed, 53 insertions(+), 32 deletions(-) diff --git a/internal/auth/kiro/aws_auth.go b/internal/auth/kiro/aws_auth.go index 53c77a8b..d082f274 100644 --- a/internal/auth/kiro/aws_auth.go +++ b/internal/auth/kiro/aws_auth.go @@ -280,6 +280,11 @@ func (k *KiroAuth) CreateTokenStorage(tokenData *KiroTokenData) *KiroTokenStorag AuthMethod: tokenData.AuthMethod, Provider: tokenData.Provider, LastRefresh: time.Now().Format(time.RFC3339), + ClientID: tokenData.ClientID, + ClientSecret: tokenData.ClientSecret, + Region: tokenData.Region, + StartURL: tokenData.StartURL, + Email: tokenData.Email, } } @@ -311,4 +316,19 @@ func (k *KiroAuth) UpdateTokenStorage(storage *KiroTokenStorage, tokenData *Kiro storage.AuthMethod = tokenData.AuthMethod storage.Provider = tokenData.Provider storage.LastRefresh = time.Now().Format(time.RFC3339) + if tokenData.ClientID != "" { + storage.ClientID = tokenData.ClientID + } + if tokenData.ClientSecret != "" { + storage.ClientSecret = tokenData.ClientSecret + } + if tokenData.Region != "" { + storage.Region = tokenData.Region + } + if tokenData.StartURL != "" { + storage.StartURL = tokenData.StartURL + } + if tokenData.Email != "" { + storage.Email = tokenData.Email + } } diff --git a/internal/auth/kiro/oauth_web.go b/internal/auth/kiro/oauth_web.go index 81c24393..6e4269c5 100644 --- a/internal/auth/kiro/oauth_web.go +++ b/internal/auth/kiro/oauth_web.go @@ -377,17 +377,18 @@ func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSess email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken) tokenData := &KiroTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - ProfileArn: profileArn, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: session.authMethod, - Provider: "AWS", - ClientID: session.clientID, - ClientSecret: session.clientSecret, - Email: email, - Region: session.region, - } + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: profileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: session.authMethod, + Provider: "AWS", + ClientID: session.clientID, + ClientSecret: session.clientSecret, + Email: email, + Region: session.region, + StartURL: session.startURL, + } h.mu.Lock() session.status = statusSuccess @@ -828,7 +829,7 @@ func (h *OAuthWebHandler) handleImportToken(c *gin.Context) { // handleManualRefresh handles manual token refresh requests from the web UI. // This allows users to trigger a token refresh when needed, without waiting -// for the automatic 5-second check and 10-minute-before-expiry refresh cycle. +// for the automatic 30-second check and 20-minute-before-expiry refresh cycle. // Uses the same refresh logic as kiro_executor.Refresh for consistency. func (h *OAuthWebHandler) handleManualRefresh(c *gin.Context) { authDir := "" diff --git a/internal/runtime/executor/kiro_executor.go b/internal/runtime/executor/kiro_executor.go index b0c14c61..cab4bcd6 100644 --- a/internal/runtime/executor/kiro_executor.go +++ b/internal/runtime/executor/kiro_executor.go @@ -3513,14 +3513,14 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c // Also check if expires_at is now in the future with sufficient buffer if expiresAt, ok := auth.Metadata["expires_at"].(string); ok { if expTime, err := time.Parse(time.RFC3339, expiresAt); err == nil { - // If token expires more than 5 minutes from now, it's still valid - if time.Until(expTime) > 5*time.Minute { + // If token expires more than 20 minutes from now, it's still valid + if time.Until(expTime) > 20*time.Minute { log.Debugf("kiro executor: token is still valid (expires in %v), skipping refresh", time.Until(expTime)) // CRITICAL FIX: Set NextRefreshAfter to prevent frequent refresh checks - // Without this, shouldRefresh() will return true again in 5 seconds + // Without this, shouldRefresh() will return true again in 30 seconds updated := auth.Clone() - // Set next refresh to 5 minutes before expiry, or at least 30 seconds from now - nextRefresh := expTime.Add(-5 * time.Minute) + // Set next refresh to 20 minutes before expiry, or at least 30 seconds from now + nextRefresh := expTime.Add(-20 * time.Minute) minNextRefresh := time.Now().Add(30 * time.Second) if nextRefresh.Before(minNextRefresh) { nextRefresh = minNextRefresh @@ -3626,9 +3626,9 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c updated.Attributes["profile_arn"] = tokenData.ProfileArn } - // NextRefreshAfter is aligned with RefreshLead (5min) + // NextRefreshAfter is aligned with RefreshLead (20min) if expiresAt, parseErr := time.Parse(time.RFC3339, tokenData.ExpiresAt); parseErr == nil { - updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute) + updated.NextRefreshAfter = expiresAt.Add(-20 * time.Minute) } log.Infof("kiro executor: token refreshed successfully, expires at %s", tokenData.ExpiresAt) diff --git a/sdk/auth/filestore.go b/sdk/auth/filestore.go index 0010be7d..9a288b10 100644 --- a/sdk/auth/filestore.go +++ b/sdk/auth/filestore.go @@ -217,11 +217,11 @@ func (s *FileTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Auth, } id := s.idFor(path, baseDir) - // Calculate NextRefreshAfter from expires_at (10 minutes before expiry) + // Calculate NextRefreshAfter from expires_at (20 minutes before expiry) var nextRefreshAfter time.Time if expiresAtStr, ok := metadata["expires_at"].(string); ok && expiresAtStr != "" { if expiresAt, err := time.Parse(time.RFC3339, expiresAtStr); err == nil { - nextRefreshAfter = expiresAt.Add(-10 * time.Minute) + nextRefreshAfter = expiresAt.Add(-20 * time.Minute) } } diff --git a/sdk/auth/kiro.go b/sdk/auth/kiro.go index b0687eba..f66be461 100644 --- a/sdk/auth/kiro.go +++ b/sdk/auth/kiro.go @@ -52,9 +52,9 @@ func (a *KiroAuthenticator) Provider() string { } // RefreshLead indicates how soon before expiry a refresh should be attempted. -// Set to 10 minutes for proactive refresh before token expiry. +// Set to 20 minutes for proactive refresh before token expiry. func (a *KiroAuthenticator) RefreshLead() *time.Duration { - d := 10 * time.Minute + d := 20 * time.Minute return &d } @@ -132,8 +132,8 @@ func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, UpdatedAt: now, Metadata: metadata, Attributes: attributes, - // NextRefreshAfter: 10 minutes before expiry - NextRefreshAfter: expiresAt.Add(-10 * time.Minute), + // NextRefreshAfter: 20 minutes before expiry + NextRefreshAfter: expiresAt.Add(-20 * time.Minute), } if tokenData.Email != "" { @@ -214,8 +214,8 @@ func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.C "source": "aws-builder-id-authcode", "email": tokenData.Email, }, - // NextRefreshAfter: 10 minutes before expiry - NextRefreshAfter: expiresAt.Add(-10 * time.Minute), + // NextRefreshAfter: 20 minutes before expiry + NextRefreshAfter: expiresAt.Add(-20 * time.Minute), } if tokenData.Email != "" { @@ -298,8 +298,8 @@ func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.C "email": tokenData.Email, "region": tokenData.Region, }, - // NextRefreshAfter: 10 minutes before expiry - NextRefreshAfter: expiresAt.Add(-10 * time.Minute), + // NextRefreshAfter: 20 minutes before expiry + NextRefreshAfter: expiresAt.Add(-20 * time.Minute), } // Display the email if extracted @@ -367,8 +367,8 @@ func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, aut updated.Metadata["refresh_token"] = tokenData.RefreshToken updated.Metadata["expires_at"] = tokenData.ExpiresAt updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization - // NextRefreshAfter: 10 minutes before expiry - updated.NextRefreshAfter = expiresAt.Add(-10 * time.Minute) + // NextRefreshAfter: 20 minutes before expiry + updated.NextRefreshAfter = expiresAt.Add(-20 * time.Minute) return updated, nil } diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 83769198..5f553bdd 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -47,7 +47,7 @@ type RefreshEvaluator interface { } const ( - refreshCheckInterval = 5 * time.Second + refreshCheckInterval = 30 * time.Second refreshPendingBackoff = time.Minute refreshFailureBackoff = 1 * time.Minute quotaBackoffBase = time.Second From 194f66ca9c1c7abfbed0c1e8874b5ed6d9ba9ec3 Mon Sep 17 00:00:00 2001 From: "yuechenglong.5" Date: Wed, 21 Jan 2026 11:03:07 +0800 Subject: [PATCH 09/10] =?UTF-8?q?feat(kiro):=20=E6=B7=BB=E5=8A=A0=E5=90=8E?= =?UTF-8?q?=E5=8F=B0=E4=BB=A4=E7=89=8C=E5=88=B7=E6=96=B0=E9=80=9A=E7=9F=A5?= =?UTF-8?q?=E6=9C=BA=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 在 BackgroundRefresher 中添加 onTokenRefreshed 回调函数和并发安全锁 - 实现 WithOnTokenRefreshed 选项函数用于设置刷新成功回调 - 在 RefreshManager 中添加 SetOnTokenRefreshed 方法支持运行时更新回调 - 为 KiroExecutor 添加 reloadAuthFromFile 方法实现文件重新加载回退机制 - 在 Watcher 中实现 NotifyTokenRefreshed 方法处理刷新通知并更新内存Auth对象 - 通过 Service.GetWatcher 连接刷新器回调到 Watcher 通知链路 - 添加方案A和方案B双重保障解决后台刷新与内存对象时间差问题 --- internal/auth/kiro/background_refresh.go | 48 +++++- internal/auth/kiro/refresh_manager.go | 50 ++++-- internal/runtime/executor/kiro_executor.go | 183 ++++++++++++++++++--- internal/watcher/watcher.go | 108 ++++++++++++ sdk/cliproxy/service.go | 22 +++ sdk/cliproxy/types.go | 14 ++ sdk/cliproxy/watcher.go | 3 + 7 files changed, 386 insertions(+), 42 deletions(-) diff --git a/internal/auth/kiro/background_refresh.go b/internal/auth/kiro/background_refresh.go index 3fecc417..1203ff47 100644 --- a/internal/auth/kiro/background_refresh.go +++ b/internal/auth/kiro/background_refresh.go @@ -50,14 +50,16 @@ func WithConcurrency(concurrency int) RefresherOption { } type BackgroundRefresher struct { - interval time.Duration - batchSize int - concurrency int - tokenRepo TokenRepository - stopCh chan struct{} - wg sync.WaitGroup - oauth *KiroOAuth - ssoClient *SSOOIDCClient + interval time.Duration + batchSize int + concurrency int + tokenRepo TokenRepository + stopCh chan struct{} + wg sync.WaitGroup + oauth *KiroOAuth + ssoClient *SSOOIDCClient + callbackMu sync.RWMutex // 保护回调函数的并发访问 + onTokenRefreshed func(tokenID string, tokenData *KiroTokenData) // 刷新成功回调 } func NewBackgroundRefresher(repo TokenRepository, opts ...RefresherOption) *BackgroundRefresher { @@ -84,6 +86,17 @@ func WithConfig(cfg *config.Config) RefresherOption { } } +// WithOnTokenRefreshed sets the callback function to be called when a token is successfully refreshed. +// The callback receives the token ID (filename) and the new token data. +// This allows external components (e.g., Watcher) to be notified of token updates. +func WithOnTokenRefreshed(callback func(tokenID string, tokenData *KiroTokenData)) RefresherOption { + return func(r *BackgroundRefresher) { + r.callbackMu.Lock() + r.onTokenRefreshed = callback + r.callbackMu.Unlock() + } +} + func (r *BackgroundRefresher) Start(ctx context.Context) { r.wg.Add(1) go func() { @@ -188,5 +201,24 @@ func (r *BackgroundRefresher) refreshSingle(ctx context.Context, token *Token) { if err := r.tokenRepo.UpdateToken(token); err != nil { log.Printf("failed to update token %s: %v", token.ID, err) + return + } + + // 方案 A: 刷新成功后触发回调，通知 Watcher 更新内存中的 Auth 对象 + r.callbackMu.RLock() + callback := r.onTokenRefreshed + r.callbackMu.RUnlock() + + if callback != nil { + // 使用 defer recover 隔离回调 panic，防止崩溃整个进程 + func() { + defer func() { + if rec := recover(); rec != nil { + log.Printf("background refresh: callback panic for token %s: %v", token.ID, rec) + } + }() + log.Printf("background refresh: notifying token refresh callback for %s", token.ID) + callback(token.ID, newTokenData) + }() } } diff --git a/internal/auth/kiro/refresh_manager.go b/internal/auth/kiro/refresh_manager.go index cd27b432..05e27a54 100644 --- a/internal/auth/kiro/refresh_manager.go +++ b/internal/auth/kiro/refresh_manager.go @@ -11,11 +11,12 @@ import ( // RefreshManager 是后台刷新器的单例管理器 type RefreshManager struct { - mu sync.Mutex - refresher *BackgroundRefresher - ctx context.Context - cancel context.CancelFunc - started bool + mu sync.Mutex + refresher *BackgroundRefresher + ctx context.Context + cancel context.CancelFunc + started bool + onTokenRefreshed func(tokenID string, tokenData *KiroTokenData) // 刷新成功回调 } var ( @@ -52,13 +53,19 @@ func (m *RefreshManager) Initialize(baseDir string, cfg *config.Config) error { repo := NewFileTokenRepository(baseDir) // 创建后台刷新器，配置参数 - m.refresher = NewBackgroundRefresher( - repo, - WithInterval(time.Minute), // 每分钟检查一次 - WithBatchSize(50), // 每批最多处理 50 个 token - WithConcurrency(10), // 最多 10 个并发刷新 - WithConfig(cfg), // 设置 OAuth 和 SSO 客户端 - ) + opts := []RefresherOption{ + WithInterval(time.Minute), // 每分钟检查一次 + WithBatchSize(50), // 每批最多处理 50 个 token + WithConcurrency(10), // 最多 10 个并发刷新 + WithConfig(cfg), // 设置 OAuth 和 SSO 客户端 + } + + // 如果已设置回调，传递给 BackgroundRefresher + if m.onTokenRefreshed != nil { + opts = append(opts, WithOnTokenRefreshed(m.onTokenRefreshed)) + } + + m.refresher = NewBackgroundRefresher(repo, opts...) log.Infof("refresh manager: initialized with base directory %s", baseDir) return nil @@ -127,6 +134,25 @@ func (m *RefreshManager) UpdateBaseDir(baseDir string) { } } +// SetOnTokenRefreshed 设置 token 刷新成功后的回调函数 +// 可以在任何时候调用，支持运行时更新回调 +// callback: 回调函数，接收 tokenID（文件名）和新的 token 数据 +func (m *RefreshManager) SetOnTokenRefreshed(callback func(tokenID string, tokenData *KiroTokenData)) { + m.mu.Lock() + defer m.mu.Unlock() + + m.onTokenRefreshed = callback + + // 如果 refresher 已经创建，使用并发安全的方式更新它的回调 + if m.refresher != nil { + m.refresher.callbackMu.Lock() + m.refresher.onTokenRefreshed = callback + m.refresher.callbackMu.Unlock() + } + + log.Debug("refresh manager: token refresh callback registered") +} + // InitializeAndStart 初始化并启动后台刷新（便捷方法） func InitializeAndStart(baseDir string, cfg *config.Config) { manager := GetRefreshManager() diff --git a/internal/runtime/executor/kiro_executor.go b/internal/runtime/executor/kiro_executor.go index 4506601d..ed6014a2 100644 --- a/internal/runtime/executor/kiro_executor.go +++ b/internal/runtime/executor/kiro_executor.go @@ -581,18 +581,30 @@ func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req // Check if token is expired before making request if e.isTokenExpired(accessToken) { - log.Infof("kiro: access token expired, attempting refresh before request") - refreshedAuth, refreshErr := e.Refresh(ctx, auth) - if refreshErr != nil { - log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr) - } else if refreshedAuth != nil { - auth = refreshedAuth - // Persist the refreshed auth to file so subsequent requests use it - if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { - log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) - } + log.Infof("kiro: access token expired, attempting recovery") + + // 方案 B: 先尝试从文件重新加载 token（后台刷新器可能已更新文件） + reloadedAuth, reloadErr := e.reloadAuthFromFile(auth) + if reloadErr == nil && reloadedAuth != nil { + // 文件中有更新的 token，使用它 + auth = reloadedAuth accessToken, profileArn = kiroCredentials(auth) - log.Infof("kiro: token refreshed successfully before request") + log.Infof("kiro: recovered token from file (background refresh), expires_at: %v", auth.Metadata["expires_at"]) + } else { + // 文件中的 token 也过期了，执行主动刷新 + log.Debugf("kiro: file reload failed (%v), attempting active refresh", reloadErr) + refreshedAuth, refreshErr := e.Refresh(ctx, auth) + if refreshErr != nil { + log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr) + } else if refreshedAuth != nil { + auth = refreshedAuth + // Persist the refreshed auth to file so subsequent requests use it + if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { + log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) + } + accessToken, profileArn = kiroCredentials(auth) + log.Infof("kiro: token refreshed successfully before request") + } } } @@ -979,18 +991,30 @@ func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut // Check if token is expired before making request if e.isTokenExpired(accessToken) { - log.Infof("kiro: access token expired, attempting refresh before stream request") - refreshedAuth, refreshErr := e.Refresh(ctx, auth) - if refreshErr != nil { - log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr) - } else if refreshedAuth != nil { - auth = refreshedAuth - // Persist the refreshed auth to file so subsequent requests use it - if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { - log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) - } + log.Infof("kiro: access token expired, attempting recovery before stream request") + + // 方案 B: 先尝试从文件重新加载 token（后台刷新器可能已更新文件） + reloadedAuth, reloadErr := e.reloadAuthFromFile(auth) + if reloadErr == nil && reloadedAuth != nil { + // 文件中有更新的 token，使用它 + auth = reloadedAuth accessToken, profileArn = kiroCredentials(auth) - log.Infof("kiro: token refreshed successfully before stream request") + log.Infof("kiro: recovered token from file (background refresh) for stream, expires_at: %v", auth.Metadata["expires_at"]) + } else { + // 文件中的 token 也过期了，执行主动刷新 + log.Debugf("kiro: file reload failed (%v), attempting active refresh for stream", reloadErr) + refreshedAuth, refreshErr := e.Refresh(ctx, auth) + if refreshErr != nil { + log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr) + } else if refreshedAuth != nil { + auth = refreshedAuth + // Persist the refreshed auth to file so subsequent requests use it + if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { + log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) + } + accessToken, profileArn = kiroCredentials(auth) + log.Infof("kiro: token refreshed successfully before stream request") + } } } @@ -3689,6 +3713,121 @@ func (e *KiroExecutor) persistRefreshedAuth(auth *cliproxyauth.Auth) error { return nil } +// reloadAuthFromFile 从文件重新加载 auth 数据（方案 B: Fallback 机制） +// 当内存中的 token 已过期时，尝试从文件读取最新的 token +// 这解决了后台刷新器已更新文件但内存中 Auth 对象尚未同步的时间差问题 +func (e *KiroExecutor) reloadAuthFromFile(auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { + if auth == nil { + return nil, fmt.Errorf("kiro executor: cannot reload nil auth") + } + + // 确定文件路径 + var authPath string + if auth.Attributes != nil { + if p := strings.TrimSpace(auth.Attributes["path"]); p != "" { + authPath = p + } + } + if authPath == "" { + fileName := strings.TrimSpace(auth.FileName) + if fileName == "" { + return nil, fmt.Errorf("kiro executor: auth has no file path or filename for reload") + } + if filepath.IsAbs(fileName) { + authPath = fileName + } else if e.cfg != nil && e.cfg.AuthDir != "" { + authPath = filepath.Join(e.cfg.AuthDir, fileName) + } else { + return nil, fmt.Errorf("kiro executor: cannot determine auth file path for reload") + } + } + + // 读取文件 + raw, err := os.ReadFile(authPath) + if err != nil { + return nil, fmt.Errorf("kiro executor: failed to read auth file %s: %w", authPath, err) + } + + // 解析 JSON + var metadata map[string]any + if err := json.Unmarshal(raw, &metadata); err != nil { + return nil, fmt.Errorf("kiro executor: failed to parse auth file %s: %w", authPath, err) + } + + // 检查文件中的 token 是否比内存中的更新 + fileExpiresAt, _ := metadata["expires_at"].(string) + fileAccessToken, _ := metadata["access_token"].(string) + memExpiresAt, _ := auth.Metadata["expires_at"].(string) + memAccessToken, _ := auth.Metadata["access_token"].(string) + + // 文件中必须有有效的 access_token + if fileAccessToken == "" { + return nil, fmt.Errorf("kiro executor: auth file has no access_token field") + } + + // 如果有 expires_at，检查是否过期 + if fileExpiresAt != "" { + fileExpTime, parseErr := time.Parse(time.RFC3339, fileExpiresAt) + if parseErr == nil { + // 如果文件中的 token 也已过期，不使用它 + if time.Now().After(fileExpTime) { + log.Debugf("kiro executor: file token also expired at %s, not using", fileExpiresAt) + return nil, fmt.Errorf("kiro executor: file token also expired") + } + } + } + + // 判断文件中的 token 是否比内存中的更新 + // 条件1: access_token 不同（说明已刷新） + // 条件2: expires_at 更新（说明已刷新） + isNewer := false + + // 优先检查 access_token 是否变化 + if fileAccessToken != memAccessToken { + isNewer = true + log.Debugf("kiro executor: file access_token differs from memory, using file token") + } + + // 如果 access_token 相同，检查 expires_at + if !isNewer && fileExpiresAt != "" && memExpiresAt != "" { + fileExpTime, fileParseErr := time.Parse(time.RFC3339, fileExpiresAt) + memExpTime, memParseErr := time.Parse(time.RFC3339, memExpiresAt) + if fileParseErr == nil && memParseErr == nil && fileExpTime.After(memExpTime) { + isNewer = true + log.Debugf("kiro executor: file expires_at (%s) is newer than memory (%s)", fileExpiresAt, memExpiresAt) + } + } + + // 如果文件中没有 expires_at 但 access_token 相同，无法判断是否更新 + if !isNewer && fileExpiresAt == "" && fileAccessToken == memAccessToken { + return nil, fmt.Errorf("kiro executor: cannot determine if file token is newer (no expires_at, same access_token)") + } + + if !isNewer { + log.Debugf("kiro executor: file token not newer than memory token") + return nil, fmt.Errorf("kiro executor: file token not newer") + } + + // 创建更新后的 auth 对象 + updated := auth.Clone() + updated.Metadata = metadata + updated.UpdatedAt = time.Now() + + // 同步更新 Attributes + if updated.Attributes == nil { + updated.Attributes = make(map[string]string) + } + if accessToken, ok := metadata["access_token"].(string); ok { + updated.Attributes["access_token"] = accessToken + } + if profileArn, ok := metadata["profile_arn"].(string); ok { + updated.Attributes["profile_arn"] = profileArn + } + + log.Infof("kiro executor: reloaded auth from file %s, new expires_at: %s", authPath, fileExpiresAt) + return updated, nil +} + // isTokenExpired checks if a JWT access token has expired. // Returns true if the token is expired or cannot be parsed. func (e *KiroExecutor) isTokenExpired(accessToken string) bool { diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 77006cf8..8141ca07 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -145,3 +145,111 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth { w.clientsMutex.RUnlock() return snapshotCoreAuths(cfg, w.authDir) } + +// NotifyTokenRefreshed 处理后台刷新器的 token 更新通知 +// 当后台刷新器成功刷新 token 后调用此方法，更新内存中的 Auth 对象 +// tokenID: token 文件名（如 kiro-xxx.json） +// accessToken: 新的 access token +// refreshToken: 新的 refresh token +// expiresAt: 新的过期时间 +func (w *Watcher) NotifyTokenRefreshed(tokenID, accessToken, refreshToken, expiresAt string) { + if w == nil { + return + } + + w.clientsMutex.Lock() + defer w.clientsMutex.Unlock() + + // 遍历 currentAuths，找到匹配的 Auth 并更新 + updated := false + for id, auth := range w.currentAuths { + if auth == nil || auth.Metadata == nil { + continue + } + + // 检查是否是 kiro 类型的 auth + authType, _ := auth.Metadata["type"].(string) + if authType != "kiro" { + continue + } + + // 多种匹配方式，解决不同来源的 auth 对象字段差异 + matched := false + + // 1. 通过 auth.ID 匹配（ID 可能包含文件名） + if !matched && auth.ID != "" { + if auth.ID == tokenID || strings.HasSuffix(auth.ID, "/"+tokenID) || strings.HasSuffix(auth.ID, "\\"+tokenID) { + matched = true + } + // ID 可能是 "kiro-xxx" 格式（无扩展名），tokenID 是 "kiro-xxx.json" + if !matched && strings.TrimSuffix(tokenID, ".json") == auth.ID { + matched = true + } + } + + // 2. 通过 auth.Attributes["path"] 匹配 + if !matched && auth.Attributes != nil { + if authPath := auth.Attributes["path"]; authPath != "" { + // 提取文件名部分进行比较 + pathBase := authPath + if idx := strings.LastIndexAny(authPath, "/\\"); idx >= 0 { + pathBase = authPath[idx+1:] + } + if pathBase == tokenID || strings.TrimSuffix(pathBase, ".json") == strings.TrimSuffix(tokenID, ".json") { + matched = true + } + } + } + + // 3. 通过 auth.FileName 匹配（原有逻辑） + if !matched && auth.FileName != "" { + if auth.FileName == tokenID || strings.HasSuffix(auth.FileName, "/"+tokenID) || strings.HasSuffix(auth.FileName, "\\"+tokenID) { + matched = true + } + } + + if matched { + // 更新内存中的 token + auth.Metadata["access_token"] = accessToken + auth.Metadata["refresh_token"] = refreshToken + auth.Metadata["expires_at"] = expiresAt + auth.Metadata["last_refresh"] = time.Now().Format(time.RFC3339) + auth.UpdatedAt = time.Now() + auth.LastRefreshedAt = time.Now() + + log.Infof("watcher: updated in-memory auth for token %s (auth ID: %s)", tokenID, id) + updated = true + + // 同时更新 runtimeAuths 中的副本（如果存在） + if w.runtimeAuths != nil { + if runtimeAuth, ok := w.runtimeAuths[id]; ok && runtimeAuth != nil { + if runtimeAuth.Metadata == nil { + runtimeAuth.Metadata = make(map[string]any) + } + runtimeAuth.Metadata["access_token"] = accessToken + runtimeAuth.Metadata["refresh_token"] = refreshToken + runtimeAuth.Metadata["expires_at"] = expiresAt + runtimeAuth.Metadata["last_refresh"] = time.Now().Format(time.RFC3339) + runtimeAuth.UpdatedAt = time.Now() + runtimeAuth.LastRefreshedAt = time.Now() + } + } + + // 发送更新通知到 authQueue + if w.authQueue != nil { + go func(authClone *coreauth.Auth) { + update := AuthUpdate{ + Action: AuthUpdateActionModify, + ID: authClone.ID, + Auth: authClone, + } + w.dispatchAuthUpdates([]AuthUpdate{update}) + }(auth.Clone()) + } + } + } + + if !updated { + log.Debugf("watcher: no matching auth found for token %s, will be picked up on next file scan", tokenID) + } +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 885304ad..750eb885 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -98,6 +98,16 @@ func (s *Service) RegisterUsagePlugin(plugin usage.Plugin) { usage.RegisterPlugin(plugin) } +// GetWatcher returns the underlying WatcherWrapper instance. +// This allows external components (e.g., RefreshManager) to interact with the watcher. +// Returns nil if the service or watcher is not initialized. +func (s *Service) GetWatcher() *WatcherWrapper { + if s == nil { + return nil + } + return s.watcher +} + // newDefaultAuthManager creates a default authentication manager with all supported providers. func newDefaultAuthManager() *sdkAuth.Manager { return sdkAuth.NewManager( @@ -575,6 +585,18 @@ func (s *Service) Run(ctx context.Context) error { } watcherWrapper.SetConfig(s.cfg) + // 方案 A: 连接 Kiro 后台刷新器回调到 Watcher + // 当后台刷新器成功刷新 token 后，立即通知 Watcher 更新内存中的 Auth 对象 + // 这解决了后台刷新与内存 Auth 对象之间的时间差问题 + kiroauth.GetRefreshManager().SetOnTokenRefreshed(func(tokenID string, tokenData *kiroauth.KiroTokenData) { + if tokenData == nil || watcherWrapper == nil { + return + } + log.Debugf("kiro refresh callback: notifying watcher for token %s", tokenID) + watcherWrapper.NotifyTokenRefreshed(tokenID, tokenData.AccessToken, tokenData.RefreshToken, tokenData.ExpiresAt) + }) + log.Debug("kiro: connected background refresh callback to watcher") + watcherCtx, watcherCancel := context.WithCancel(context.Background()) s.watcherCancel = watcherCancel if err = watcherWrapper.Start(watcherCtx); err != nil { diff --git a/sdk/cliproxy/types.go b/sdk/cliproxy/types.go index 1521dffe..ee8f761d 100644 --- a/sdk/cliproxy/types.go +++ b/sdk/cliproxy/types.go @@ -89,6 +89,7 @@ type WatcherWrapper struct { snapshotAuths func() []*coreauth.Auth setUpdateQueue func(queue chan<- watcher.AuthUpdate) dispatchRuntimeUpdate func(update watcher.AuthUpdate) bool + notifyTokenRefreshed func(tokenID, accessToken, refreshToken, expiresAt string) // 方案 A: 后台刷新通知 } // Start proxies to the underlying watcher Start implementation. @@ -146,3 +147,16 @@ func (w *WatcherWrapper) SetAuthUpdateQueue(queue chan<- watcher.AuthUpdate) { } w.setUpdateQueue(queue) } + +// NotifyTokenRefreshed 通知 Watcher 后台刷新器已更新 token +// 这是方案 A 的核心方法，用于解决后台刷新与内存 Auth 对象的时间差问题 +// tokenID: token 文件名（如 kiro-xxx.json） +// accessToken: 新的 access token +// refreshToken: 新的 refresh token +// expiresAt: 新的过期时间（RFC3339 格式） +func (w *WatcherWrapper) NotifyTokenRefreshed(tokenID, accessToken, refreshToken, expiresAt string) { + if w == nil || w.notifyTokenRefreshed == nil { + return + } + w.notifyTokenRefreshed(tokenID, accessToken, refreshToken, expiresAt) +} diff --git a/sdk/cliproxy/watcher.go b/sdk/cliproxy/watcher.go index caeadf19..e6e91bdd 100644 --- a/sdk/cliproxy/watcher.go +++ b/sdk/cliproxy/watcher.go @@ -31,5 +31,8 @@ func defaultWatcherFactory(configPath, authDir string, reload func(*config.Confi dispatchRuntimeUpdate: func(update watcher.AuthUpdate) bool { return w.DispatchRuntimeAuthUpdate(update) }, + notifyTokenRefreshed: func(tokenID, accessToken, refreshToken, expiresAt string) { + w.NotifyTokenRefreshed(tokenID, accessToken, refreshToken, expiresAt) + }, }, nil } From 4c8026ac3dd235f7c59857ecfd98ae0504d5e0bb Mon Sep 17 00:00:00 2001 From: "yuechenglong.5" Date: Wed, 21 Jan 2026 21:38:47 +0800 Subject: [PATCH 10/10] =?UTF-8?q?chore(build):=20=E6=9B=B4=E6=96=B0=20.git?= =?UTF-8?q?ignore=20=E6=96=87=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 添加 *.bak 文件扩展名到忽略列表 --- .gitignore | 1 + internal/auth/kiro/aws.go.bak | 305 ---- internal/auth/kiro/oauth_web.go.bak | 385 ----- internal/auth/kiro/sso_oidc.go.bak | 1371 ----------------- .../translator/kiro/common/utf8_stream.go | 97 -- .../kiro/common/utf8_stream_test.go | 402 ----- sdk/auth/kiro.go.bak | 470 ------ 7 files changed, 1 insertion(+), 3030 deletions(-) delete mode 100644 internal/auth/kiro/aws.go.bak delete mode 100644 internal/auth/kiro/oauth_web.go.bak delete mode 100644 internal/auth/kiro/sso_oidc.go.bak delete mode 100644 internal/translator/kiro/common/utf8_stream.go delete mode 100644 internal/translator/kiro/common/utf8_stream_test.go delete mode 100644 sdk/auth/kiro.go.bak diff --git a/.gitignore b/.gitignore index 29cf765b..bab49132 100644 --- a/.gitignore +++ b/.gitignore @@ -50,3 +50,4 @@ _bmad-output/* # macOS .DS_Store ._* +*.bak diff --git a/internal/auth/kiro/aws.go.bak b/internal/auth/kiro/aws.go.bak deleted file mode 100644 index ba73af4d..00000000 --- a/internal/auth/kiro/aws.go.bak +++ /dev/null @@ -1,305 +0,0 @@ -// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API. -// It includes interfaces and implementations for token storage and authentication methods. -package kiro - -import ( - "encoding/base64" - "encoding/json" - "fmt" - "os" - "path/filepath" - "strings" -) - -// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow -type PKCECodes struct { - // CodeVerifier is the cryptographically random string used to correlate - // the authorization request to the token request - CodeVerifier string `json:"code_verifier"` - // CodeChallenge is the SHA256 hash of the code verifier, base64url-encoded - CodeChallenge string `json:"code_challenge"` -} - -// KiroTokenData holds OAuth token information from AWS CodeWhisperer (Kiro) -type KiroTokenData struct { - // AccessToken is the OAuth2 access token for API access - AccessToken string `json:"accessToken"` - // RefreshToken is used to obtain new access tokens - RefreshToken string `json:"refreshToken"` - // ProfileArn is the AWS CodeWhisperer profile ARN - ProfileArn string `json:"profileArn"` - // ExpiresAt is the timestamp when the token expires - ExpiresAt string `json:"expiresAt"` - // AuthMethod indicates the authentication method used (e.g., "builder-id", "social") - AuthMethod string `json:"authMethod"` - // Provider indicates the OAuth provider (e.g., "AWS", "Google") - Provider string `json:"provider"` - // ClientID is the OIDC client ID (needed for token refresh) - ClientID string `json:"clientId,omitempty"` - // ClientSecret is the OIDC client secret (needed for token refresh) - ClientSecret string `json:"clientSecret,omitempty"` - // Email is the user's email address (used for file naming) - Email string `json:"email,omitempty"` - // StartURL is the IDC/Identity Center start URL (only for IDC auth method) - StartURL string `json:"startUrl,omitempty"` - // Region is the AWS region for IDC authentication (only for IDC auth method) - Region string `json:"region,omitempty"` -} - -// KiroAuthBundle aggregates authentication data after OAuth flow completion -type KiroAuthBundle struct { - // TokenData contains the OAuth tokens from the authentication flow - TokenData KiroTokenData `json:"token_data"` - // LastRefresh is the timestamp of the last token refresh - LastRefresh string `json:"last_refresh"` -} - -// KiroUsageInfo represents usage information from CodeWhisperer API -type KiroUsageInfo struct { - // SubscriptionTitle is the subscription plan name (e.g., "KIRO FREE") - SubscriptionTitle string `json:"subscription_title"` - // CurrentUsage is the current credit usage - CurrentUsage float64 `json:"current_usage"` - // UsageLimit is the maximum credit limit - UsageLimit float64 `json:"usage_limit"` - // NextReset is the timestamp of the next usage reset - NextReset string `json:"next_reset"` -} - -// KiroModel represents a model available through the CodeWhisperer API -type KiroModel struct { - // ModelID is the unique identifier for the model - ModelID string `json:"modelId"` - // ModelName is the human-readable name - ModelName string `json:"modelName"` - // Description is the model description - Description string `json:"description"` - // RateMultiplier is the credit multiplier for this model - RateMultiplier float64 `json:"rateMultiplier"` - // RateUnit is the unit for rate calculation (e.g., "credit") - RateUnit string `json:"rateUnit"` - // MaxInputTokens is the maximum input token limit - MaxInputTokens int `json:"maxInputTokens,omitempty"` -} - -// KiroIDETokenFile is the default path to Kiro IDE's token file -const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json" - -// LoadKiroIDEToken loads token data from Kiro IDE's token file. -func LoadKiroIDEToken() (*KiroTokenData, error) { - homeDir, err := os.UserHomeDir() - if err != nil { - return nil, fmt.Errorf("failed to get home directory: %w", err) - } - - tokenPath := filepath.Join(homeDir, KiroIDETokenFile) - data, err := os.ReadFile(tokenPath) - if err != nil { - return nil, fmt.Errorf("failed to read Kiro IDE token file (%s): %w", tokenPath, err) - } - - var token KiroTokenData - if err := json.Unmarshal(data, &token); err != nil { - return nil, fmt.Errorf("failed to parse Kiro IDE token: %w", err) - } - - if token.AccessToken == "" { - return nil, fmt.Errorf("access token is empty in Kiro IDE token file") - } - - return &token, nil -} - -// LoadKiroTokenFromPath loads token data from a custom path. -// This supports multiple accounts by allowing different token files. -func LoadKiroTokenFromPath(tokenPath string) (*KiroTokenData, error) { - // Expand ~ to home directory - if len(tokenPath) > 0 && tokenPath[0] == '~' { - homeDir, err := os.UserHomeDir() - if err != nil { - return nil, fmt.Errorf("failed to get home directory: %w", err) - } - tokenPath = filepath.Join(homeDir, tokenPath[1:]) - } - - data, err := os.ReadFile(tokenPath) - if err != nil { - return nil, fmt.Errorf("failed to read token file (%s): %w", tokenPath, err) - } - - var token KiroTokenData - if err := json.Unmarshal(data, &token); err != nil { - return nil, fmt.Errorf("failed to parse token file: %w", err) - } - - if token.AccessToken == "" { - return nil, fmt.Errorf("access token is empty in token file") - } - - return &token, nil -} - -// ListKiroTokenFiles lists all Kiro token files in the cache directory. -// This supports multiple accounts by finding all token files. -func ListKiroTokenFiles() ([]string, error) { - homeDir, err := os.UserHomeDir() - if err != nil { - return nil, fmt.Errorf("failed to get home directory: %w", err) - } - - cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache") - - // Check if directory exists - if _, err := os.Stat(cacheDir); os.IsNotExist(err) { - return nil, nil // No token files - } - - entries, err := os.ReadDir(cacheDir) - if err != nil { - return nil, fmt.Errorf("failed to read cache directory: %w", err) - } - - var tokenFiles []string - for _, entry := range entries { - if entry.IsDir() { - continue - } - name := entry.Name() - // Look for kiro token files only (avoid matching unrelated AWS SSO cache files) - if strings.HasSuffix(name, ".json") && strings.HasPrefix(name, "kiro") { - tokenFiles = append(tokenFiles, filepath.Join(cacheDir, name)) - } - } - - return tokenFiles, nil -} - -// LoadAllKiroTokens loads all Kiro tokens from the cache directory. -// This supports multiple accounts. -func LoadAllKiroTokens() ([]*KiroTokenData, error) { - files, err := ListKiroTokenFiles() - if err != nil { - return nil, err - } - - var tokens []*KiroTokenData - for _, file := range files { - token, err := LoadKiroTokenFromPath(file) - if err != nil { - // Skip invalid token files - continue - } - tokens = append(tokens, token) - } - - return tokens, nil -} - -// JWTClaims represents the claims we care about from a JWT token. -// JWT tokens from Kiro/AWS contain user information in the payload. -type JWTClaims struct { - Email string `json:"email,omitempty"` - Sub string `json:"sub,omitempty"` - PreferredUser string `json:"preferred_username,omitempty"` - Name string `json:"name,omitempty"` - Iss string `json:"iss,omitempty"` -} - -// ExtractEmailFromJWT extracts the user's email from a JWT access token. -// JWT tokens typically have format: header.payload.signature -// The payload is base64url-encoded JSON containing user claims. -func ExtractEmailFromJWT(accessToken string) string { - if accessToken == "" { - return "" - } - - // JWT format: header.payload.signature - parts := strings.Split(accessToken, ".") - if len(parts) != 3 { - return "" - } - - // Decode the payload (second part) - payload := parts[1] - - // Add padding if needed (base64url requires padding) - switch len(payload) % 4 { - case 2: - payload += "==" - case 3: - payload += "=" - } - - decoded, err := base64.URLEncoding.DecodeString(payload) - if err != nil { - // Try RawURLEncoding (no padding) - decoded, err = base64.RawURLEncoding.DecodeString(parts[1]) - if err != nil { - return "" - } - } - - var claims JWTClaims - if err := json.Unmarshal(decoded, &claims); err != nil { - return "" - } - - // Return email if available - if claims.Email != "" { - return claims.Email - } - - // Fallback to preferred_username (some providers use this) - if claims.PreferredUser != "" && strings.Contains(claims.PreferredUser, "@") { - return claims.PreferredUser - } - - // Fallback to sub if it looks like an email - if claims.Sub != "" && strings.Contains(claims.Sub, "@") { - return claims.Sub - } - - return "" -} - -// SanitizeEmailForFilename sanitizes an email address for use in a filename. -// Replaces special characters with underscores and prevents path traversal attacks. -// Also handles URL-encoded characters to prevent encoded path traversal attempts. -func SanitizeEmailForFilename(email string) string { - if email == "" { - return "" - } - - result := email - - // First, handle URL-encoded path traversal attempts (%2F, %2E, %5C, etc.) - // This prevents encoded characters from bypassing the sanitization. - // Note: We replace % last to catch any remaining encodings including double-encoding (%252F) - result = strings.ReplaceAll(result, "%2F", "_") // / - result = strings.ReplaceAll(result, "%2f", "_") - result = strings.ReplaceAll(result, "%5C", "_") // \ - result = strings.ReplaceAll(result, "%5c", "_") - result = strings.ReplaceAll(result, "%2E", "_") // . - result = strings.ReplaceAll(result, "%2e", "_") - result = strings.ReplaceAll(result, "%00", "_") // null byte - result = strings.ReplaceAll(result, "%", "_") // Catch remaining % to prevent double-encoding attacks - - // Replace characters that are problematic in filenames - // Keep @ and . in middle but replace other special characters - for _, char := range []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " ", "\x00"} { - result = strings.ReplaceAll(result, char, "_") - } - - // Prevent path traversal: replace leading dots in each path component - // This handles cases like "../../../etc/passwd" → "_.._.._.._etc_passwd" - parts := strings.Split(result, "_") - for i, part := range parts { - for strings.HasPrefix(part, ".") { - part = "_" + part[1:] - } - parts[i] = part - } - result = strings.Join(parts, "_") - - return result -} diff --git a/internal/auth/kiro/oauth_web.go.bak b/internal/auth/kiro/oauth_web.go.bak deleted file mode 100644 index 22d7809b..00000000 --- a/internal/auth/kiro/oauth_web.go.bak +++ /dev/null @@ -1,385 +0,0 @@ -// Package kiro provides OAuth Web authentication for Kiro. -package kiro - -import ( - "context" - "crypto/rand" - "encoding/base64" - "fmt" - "html/template" - "net/http" - "sync" - "time" - - "github.com/gin-gonic/gin" - "github.com/router-for-me/CLIProxyAPI/v6/internal/config" - log "github.com/sirupsen/logrus" -) - -const ( - defaultSessionExpiry = 10 * time.Minute - pollIntervalSeconds = 5 -) - -type authSessionStatus string - -const ( - statusPending authSessionStatus = "pending" - statusSuccess authSessionStatus = "success" - statusFailed authSessionStatus = "failed" -) - -type webAuthSession struct { - stateID string - deviceCode string - userCode string - authURL string - verificationURI string - expiresIn int - interval int - status authSessionStatus - startedAt time.Time - completedAt time.Time - expiresAt time.Time - error string - tokenData *KiroTokenData - ssoClient *SSOOIDCClient - clientID string - clientSecret string - region string - cancelFunc context.CancelFunc -} - -type OAuthWebHandler struct { - cfg *config.Config - sessions map[string]*webAuthSession - mu sync.RWMutex - onTokenObtained func(*KiroTokenData) -} - -func NewOAuthWebHandler(cfg *config.Config) *OAuthWebHandler { - return &OAuthWebHandler{ - cfg: cfg, - sessions: make(map[string]*webAuthSession), - } -} - -func (h *OAuthWebHandler) SetTokenCallback(callback func(*KiroTokenData)) { - h.onTokenObtained = callback -} - -func (h *OAuthWebHandler) RegisterRoutes(router gin.IRouter) { - oauth := router.Group("/v0/oauth/kiro") - { - oauth.GET("/start", h.handleStart) - oauth.GET("/callback", h.handleCallback) - oauth.GET("/status", h.handleStatus) - } -} - -func generateStateID() (string, error) { - b := make([]byte, 16) - if _, err := rand.Read(b); err != nil { - return "", err - } - return base64.RawURLEncoding.EncodeToString(b), nil -} - -func (h *OAuthWebHandler) handleStart(c *gin.Context) { - stateID, err := generateStateID() - if err != nil { - h.renderError(c, "Failed to generate state parameter") - return - } - - region := defaultIDCRegion - startURL := builderIDStartURL - - ssoClient := NewSSOOIDCClient(h.cfg) - - regResp, err := ssoClient.RegisterClientWithRegion(c.Request.Context(), region) - if err != nil { - log.Errorf("OAuth Web: failed to register client: %v", err) - h.renderError(c, fmt.Sprintf("Failed to register client: %v", err)) - return - } - - authResp, err := ssoClient.StartDeviceAuthorizationWithIDC( - c.Request.Context(), - regResp.ClientID, - regResp.ClientSecret, - startURL, - region, - ) - if err != nil { - log.Errorf("OAuth Web: failed to start device authorization: %v", err) - h.renderError(c, fmt.Sprintf("Failed to start device authorization: %v", err)) - return - } - - ctx, cancel := context.WithTimeout(context.Background(), time.Duration(authResp.ExpiresIn)*time.Second) - - session := &webAuthSession{ - stateID: stateID, - deviceCode: authResp.DeviceCode, - userCode: authResp.UserCode, - authURL: authResp.VerificationURIComplete, - verificationURI: authResp.VerificationURI, - expiresIn: authResp.ExpiresIn, - interval: authResp.Interval, - status: statusPending, - startedAt: time.Now(), - ssoClient: ssoClient, - clientID: regResp.ClientID, - clientSecret: regResp.ClientSecret, - region: region, - cancelFunc: cancel, - } - - h.mu.Lock() - h.sessions[stateID] = session - h.mu.Unlock() - - go h.pollForToken(ctx, session) - - h.renderStartPage(c, session) -} - -func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSession) { - defer session.cancelFunc() - - interval := time.Duration(session.interval) * time.Second - if interval < time.Duration(pollIntervalSeconds)*time.Second { - interval = time.Duration(pollIntervalSeconds) * time.Second - } - - ticker := time.NewTicker(interval) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - h.mu.Lock() - if session.status == statusPending { - session.status = statusFailed - session.error = "Authentication timed out" - } - h.mu.Unlock() - return - case <-ticker.C: - tokenResp, err := h.ssoClient(session).CreateTokenWithRegion( - ctx, - session.clientID, - session.clientSecret, - session.deviceCode, - session.region, - ) - - if err != nil { - errStr := err.Error() - if errStr == ErrAuthorizationPending.Error() { - continue - } - if errStr == ErrSlowDown.Error() { - interval += 5 * time.Second - ticker.Reset(interval) - continue - } - - h.mu.Lock() - session.status = statusFailed - session.error = errStr - session.completedAt = time.Now() - h.mu.Unlock() - - log.Errorf("OAuth Web: token polling failed: %v", err) - return - } - - expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) - profileArn := session.ssoClient.fetchProfileArn(ctx, tokenResp.AccessToken) - email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken) - - tokenData := &KiroTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - ProfileArn: profileArn, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: "builder-id", - Provider: "AWS", - ClientID: session.clientID, - ClientSecret: session.clientSecret, - Email: email, - } - - h.mu.Lock() - session.status = statusSuccess - session.completedAt = time.Now() - session.expiresAt = expiresAt - session.tokenData = tokenData - h.mu.Unlock() - - if h.onTokenObtained != nil { - h.onTokenObtained(tokenData) - } - - log.Infof("OAuth Web: authentication successful for %s", email) - return - } - } -} - -func (h *OAuthWebHandler) ssoClient(session *webAuthSession) *SSOOIDCClient { - return session.ssoClient -} - -func (h *OAuthWebHandler) handleCallback(c *gin.Context) { - stateID := c.Query("state") - errParam := c.Query("error") - - if errParam != "" { - h.renderError(c, errParam) - return - } - - if stateID == "" { - h.renderError(c, "Missing state parameter") - return - } - - h.mu.RLock() - session, exists := h.sessions[stateID] - h.mu.RUnlock() - - if !exists { - h.renderError(c, "Invalid or expired session") - return - } - - if session.status == statusSuccess { - h.renderSuccess(c, session) - } else if session.status == statusFailed { - h.renderError(c, session.error) - } else { - c.Redirect(http.StatusFound, "/v0/oauth/kiro/start") - } -} - -func (h *OAuthWebHandler) handleStatus(c *gin.Context) { - stateID := c.Query("state") - if stateID == "" { - c.JSON(http.StatusBadRequest, gin.H{"error": "missing state parameter"}) - return - } - - h.mu.RLock() - session, exists := h.sessions[stateID] - h.mu.RUnlock() - - if !exists { - c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) - return - } - - response := gin.H{ - "status": string(session.status), - } - - switch session.status { - case statusPending: - elapsed := time.Since(session.startedAt).Seconds() - remaining := float64(session.expiresIn) - elapsed - if remaining < 0 { - remaining = 0 - } - response["remaining_seconds"] = int(remaining) - case statusSuccess: - response["completed_at"] = session.completedAt.Format(time.RFC3339) - response["expires_at"] = session.expiresAt.Format(time.RFC3339) - case statusFailed: - response["error"] = session.error - response["failed_at"] = session.completedAt.Format(time.RFC3339) - } - - c.JSON(http.StatusOK, response) -} - -func (h *OAuthWebHandler) renderStartPage(c *gin.Context, session *webAuthSession) { - tmpl, err := template.New("start").Parse(oauthWebStartPageHTML) - if err != nil { - log.Errorf("OAuth Web: failed to parse template: %v", err) - c.String(http.StatusInternalServerError, "Template error") - return - } - - data := map[string]interface{}{ - "AuthURL": session.authURL, - "UserCode": session.userCode, - "ExpiresIn": session.expiresIn, - "StateID": session.stateID, - } - - c.Header("Content-Type", "text/html; charset=utf-8") - if err := tmpl.Execute(c.Writer, data); err != nil { - log.Errorf("OAuth Web: failed to render template: %v", err) - } -} - -func (h *OAuthWebHandler) renderError(c *gin.Context, errMsg string) { - tmpl, err := template.New("error").Parse(oauthWebErrorPageHTML) - if err != nil { - log.Errorf("OAuth Web: failed to parse error template: %v", err) - c.String(http.StatusInternalServerError, "Template error") - return - } - - data := map[string]interface{}{ - "Error": errMsg, - } - - c.Header("Content-Type", "text/html; charset=utf-8") - c.Status(http.StatusBadRequest) - if err := tmpl.Execute(c.Writer, data); err != nil { - log.Errorf("OAuth Web: failed to render error template: %v", err) - } -} - -func (h *OAuthWebHandler) renderSuccess(c *gin.Context, session *webAuthSession) { - tmpl, err := template.New("success").Parse(oauthWebSuccessPageHTML) - if err != nil { - log.Errorf("OAuth Web: failed to parse success template: %v", err) - c.String(http.StatusInternalServerError, "Template error") - return - } - - data := map[string]interface{}{ - "ExpiresAt": session.expiresAt.Format(time.RFC3339), - } - - c.Header("Content-Type", "text/html; charset=utf-8") - if err := tmpl.Execute(c.Writer, data); err != nil { - log.Errorf("OAuth Web: failed to render success template: %v", err) - } -} - -func (h *OAuthWebHandler) CleanupExpiredSessions() { - h.mu.Lock() - defer h.mu.Unlock() - - now := time.Now() - for id, session := range h.sessions { - if session.status != statusPending && now.Sub(session.completedAt) > 30*time.Minute { - delete(h.sessions, id) - } else if session.status == statusPending && now.Sub(session.startedAt) > defaultSessionExpiry { - session.cancelFunc() - delete(h.sessions, id) - } - } -} - -func (h *OAuthWebHandler) GetSession(stateID string) (*webAuthSession, bool) { - h.mu.RLock() - defer h.mu.RUnlock() - session, exists := h.sessions[stateID] - return session, exists -} diff --git a/internal/auth/kiro/sso_oidc.go.bak b/internal/auth/kiro/sso_oidc.go.bak deleted file mode 100644 index ab44e55f..00000000 --- a/internal/auth/kiro/sso_oidc.go.bak +++ /dev/null @@ -1,1371 +0,0 @@ -// Package kiro provides AWS SSO OIDC authentication for Kiro. -package kiro - -import ( - "bufio" - "context" - "crypto/rand" - "crypto/sha256" - "encoding/base64" - "encoding/json" - "errors" - "fmt" - "html" - "io" - "net" - "net/http" - "os" - "strings" - "time" - - "github.com/router-for-me/CLIProxyAPI/v6/internal/browser" - "github.com/router-for-me/CLIProxyAPI/v6/internal/config" - "github.com/router-for-me/CLIProxyAPI/v6/internal/util" - log "github.com/sirupsen/logrus" -) - -const ( - // AWS SSO OIDC endpoints - ssoOIDCEndpoint = "https://oidc.us-east-1.amazonaws.com" - - // Kiro's start URL for Builder ID - builderIDStartURL = "https://view.awsapps.com/start" - - // Default region for IDC - defaultIDCRegion = "us-east-1" - - // Polling interval - pollInterval = 5 * time.Second - - // Authorization code flow callback - authCodeCallbackPath = "/oauth/callback" - authCodeCallbackPort = 19877 - - // User-Agent to match official Kiro IDE - kiroUserAgent = "KiroIDE" - - // IDC token refresh headers (matching Kiro IDE behavior) - idcAmzUserAgent = "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE" -) - -// Sentinel errors for OIDC token polling -var ( - ErrAuthorizationPending = errors.New("authorization_pending") - ErrSlowDown = errors.New("slow_down") -) - -// SSOOIDCClient handles AWS SSO OIDC authentication. -type SSOOIDCClient struct { - httpClient *http.Client - cfg *config.Config -} - -// NewSSOOIDCClient creates a new SSO OIDC client. -func NewSSOOIDCClient(cfg *config.Config) *SSOOIDCClient { - client := &http.Client{Timeout: 30 * time.Second} - if cfg != nil { - client = util.SetProxy(&cfg.SDKConfig, client) - } - return &SSOOIDCClient{ - httpClient: client, - cfg: cfg, - } -} - -// RegisterClientResponse from AWS SSO OIDC. -type RegisterClientResponse struct { - ClientID string `json:"clientId"` - ClientSecret string `json:"clientSecret"` - ClientIDIssuedAt int64 `json:"clientIdIssuedAt"` - ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"` -} - -// StartDeviceAuthResponse from AWS SSO OIDC. -type StartDeviceAuthResponse struct { - DeviceCode string `json:"deviceCode"` - UserCode string `json:"userCode"` - VerificationURI string `json:"verificationUri"` - VerificationURIComplete string `json:"verificationUriComplete"` - ExpiresIn int `json:"expiresIn"` - Interval int `json:"interval"` -} - -// CreateTokenResponse from AWS SSO OIDC. -type CreateTokenResponse struct { - AccessToken string `json:"accessToken"` - TokenType string `json:"tokenType"` - ExpiresIn int `json:"expiresIn"` - RefreshToken string `json:"refreshToken"` -} - -// getOIDCEndpoint returns the OIDC endpoint for the given region. -func getOIDCEndpoint(region string) string { - if region == "" { - region = defaultIDCRegion - } - return fmt.Sprintf("https://oidc.%s.amazonaws.com", region) -} - -// promptInput prompts the user for input with an optional default value. -func promptInput(prompt, defaultValue string) string { - reader := bufio.NewReader(os.Stdin) - if defaultValue != "" { - fmt.Printf("%s [%s]: ", prompt, defaultValue) - } else { - fmt.Printf("%s: ", prompt) - } - input, err := reader.ReadString('\n') - if err != nil { - log.Warnf("Error reading input: %v", err) - return defaultValue - } - input = strings.TrimSpace(input) - if input == "" { - return defaultValue - } - return input -} - -// promptSelect prompts the user to select from options using number input. -func promptSelect(prompt string, options []string) int { - reader := bufio.NewReader(os.Stdin) - - for { - fmt.Println(prompt) - for i, opt := range options { - fmt.Printf(" %d) %s\n", i+1, opt) - } - fmt.Printf("Enter selection (1-%d): ", len(options)) - - input, err := reader.ReadString('\n') - if err != nil { - log.Warnf("Error reading input: %v", err) - return 0 // Default to first option on error - } - input = strings.TrimSpace(input) - - // Parse the selection - var selection int - if _, err := fmt.Sscanf(input, "%d", &selection); err != nil || selection < 1 || selection > len(options) { - fmt.Printf("Invalid selection '%s'. Please enter a number between 1 and %d.\n\n", input, len(options)) - continue - } - return selection - 1 - } -} - -// RegisterClientWithRegion registers a new OIDC client with AWS using a specific region. -func (c *SSOOIDCClient) RegisterClientWithRegion(ctx context.Context, region string) (*RegisterClientResponse, error) { - endpoint := getOIDCEndpoint(region) - - payload := map[string]interface{}{ - "clientName": "Kiro IDE", - "clientType": "public", - "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"}, - "grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"}, - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/client/register", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode) - } - - var result RegisterClientResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// StartDeviceAuthorizationWithIDC starts the device authorization flow for IDC. -func (c *SSOOIDCClient) StartDeviceAuthorizationWithIDC(ctx context.Context, clientID, clientSecret, startURL, region string) (*StartDeviceAuthResponse, error) { - endpoint := getOIDCEndpoint(region) - - payload := map[string]string{ - "clientId": clientID, - "clientSecret": clientSecret, - "startUrl": startURL, - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/device_authorization", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode) - } - - var result StartDeviceAuthResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// CreateTokenWithRegion polls for the access token after user authorization using a specific region. -func (c *SSOOIDCClient) CreateTokenWithRegion(ctx context.Context, clientID, clientSecret, deviceCode, region string) (*CreateTokenResponse, error) { - endpoint := getOIDCEndpoint(region) - - payload := map[string]string{ - "clientId": clientID, - "clientSecret": clientSecret, - "deviceCode": deviceCode, - "grantType": "urn:ietf:params:oauth:grant-type:device_code", - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - // Check for pending authorization - if resp.StatusCode == http.StatusBadRequest { - var errResp struct { - Error string `json:"error"` - } - if json.Unmarshal(respBody, &errResp) == nil { - if errResp.Error == "authorization_pending" { - return nil, ErrAuthorizationPending - } - if errResp.Error == "slow_down" { - return nil, ErrSlowDown - } - } - log.Debugf("create token failed: %s", string(respBody)) - return nil, fmt.Errorf("create token failed") - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode) - } - - var result CreateTokenResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// RefreshTokenWithRegion refreshes an access token using the refresh token with a specific region. -func (c *SSOOIDCClient) RefreshTokenWithRegion(ctx context.Context, clientID, clientSecret, refreshToken, region, startURL string) (*KiroTokenData, error) { - endpoint := getOIDCEndpoint(region) - - payload := map[string]string{ - "clientId": clientID, - "clientSecret": clientSecret, - "refreshToken": refreshToken, - "grantType": "refresh_token", - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - - // Set headers matching kiro2api's IDC token refresh - // These headers are required for successful IDC token refresh - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region)) - req.Header.Set("Connection", "keep-alive") - req.Header.Set("x-amz-user-agent", idcAmzUserAgent) - req.Header.Set("Accept", "*/*") - req.Header.Set("Accept-Language", "*") - req.Header.Set("sec-fetch-mode", "cors") - req.Header.Set("User-Agent", "node") - req.Header.Set("Accept-Encoding", "br, gzip, deflate") - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Warnf("IDC token refresh failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode) - } - - var result CreateTokenResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second) - - return &KiroTokenData{ - AccessToken: result.AccessToken, - RefreshToken: result.RefreshToken, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: "idc", - Provider: "AWS", - ClientID: clientID, - ClientSecret: clientSecret, - StartURL: startURL, - Region: region, - }, nil -} - -// LoginWithIDC performs the full device code flow for AWS Identity Center (IDC). -func (c *SSOOIDCClient) LoginWithIDC(ctx context.Context, startURL, region string) (*KiroTokenData, error) { - fmt.Println("\n╔══════════════════════════════════════════════════════════╗") - fmt.Println("║ Kiro Authentication (AWS Identity Center) ║") - fmt.Println("╚══════════════════════════════════════════════════════════╝") - - // Step 1: Register client with the specified region - fmt.Println("\nRegistering client...") - regResp, err := c.RegisterClientWithRegion(ctx, region) - if err != nil { - return nil, fmt.Errorf("failed to register client: %w", err) - } - log.Debugf("Client registered: %s", regResp.ClientID) - - // Step 2: Start device authorization with IDC start URL - fmt.Println("Starting device authorization...") - authResp, err := c.StartDeviceAuthorizationWithIDC(ctx, regResp.ClientID, regResp.ClientSecret, startURL, region) - if err != nil { - return nil, fmt.Errorf("failed to start device auth: %w", err) - } - - // Step 3: Show user the verification URL - fmt.Printf("\n") - fmt.Println("════════════════════════════════════════════════════════════") - fmt.Printf(" Confirm the following code in the browser:\n") - fmt.Printf(" Code: %s\n", authResp.UserCode) - fmt.Println("════════════════════════════════════════════════════════════") - fmt.Printf("\n Open this URL: %s\n\n", authResp.VerificationURIComplete) - - // Set incognito mode based on config - if c.cfg != nil { - browser.SetIncognitoMode(c.cfg.IncognitoBrowser) - if !c.cfg.IncognitoBrowser { - log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.") - } else { - log.Debug("kiro: using incognito mode for multi-account support") - } - } else { - browser.SetIncognitoMode(true) - log.Debug("kiro: using incognito mode for multi-account support (default)") - } - - // Open browser - if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil { - log.Warnf("Could not open browser automatically: %v", err) - fmt.Println(" Please open the URL manually in your browser.") - } else { - fmt.Println(" (Browser opened automatically)") - } - - // Step 4: Poll for token - fmt.Println("Waiting for authorization...") - - interval := pollInterval - if authResp.Interval > 0 { - interval = time.Duration(authResp.Interval) * time.Second - } - - deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second) - - for time.Now().Before(deadline) { - select { - case <-ctx.Done(): - browser.CloseBrowser() - return nil, ctx.Err() - case <-time.After(interval): - tokenResp, err := c.CreateTokenWithRegion(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode, region) - if err != nil { - if errors.Is(err, ErrAuthorizationPending) { - fmt.Print(".") - continue - } - if errors.Is(err, ErrSlowDown) { - interval += 5 * time.Second - continue - } - browser.CloseBrowser() - return nil, fmt.Errorf("token creation failed: %w", err) - } - - fmt.Println("\n\n✓ Authorization successful!") - - // Close the browser window - if err := browser.CloseBrowser(); err != nil { - log.Debugf("Failed to close browser: %v", err) - } - - // Step 5: Get profile ARN from CodeWhisperer API - fmt.Println("Fetching profile information...") - profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken) - - // Fetch user email - email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken) - if email != "" { - fmt.Printf(" Logged in as: %s\n", email) - } - - expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) - - return &KiroTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - ProfileArn: profileArn, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: "idc", - Provider: "AWS", - ClientID: regResp.ClientID, - ClientSecret: regResp.ClientSecret, - Email: email, - StartURL: startURL, - Region: region, - }, nil - } - } - - // Close browser on timeout - if err := browser.CloseBrowser(); err != nil { - log.Debugf("Failed to close browser on timeout: %v", err) - } - return nil, fmt.Errorf("authorization timed out") -} - -// LoginWithMethodSelection prompts the user to select between Builder ID and IDC, then performs the login. -func (c *SSOOIDCClient) LoginWithMethodSelection(ctx context.Context) (*KiroTokenData, error) { - fmt.Println("\n╔══════════════════════════════════════════════════════════╗") - fmt.Println("║ Kiro Authentication (AWS) ║") - fmt.Println("╚══════════════════════════════════════════════════════════╝") - - // Prompt for login method - options := []string{ - "Use with Builder ID (personal AWS account)", - "Use with IDC Account (organization SSO)", - } - selection := promptSelect("\n? Select login method:", options) - - if selection == 0 { - // Builder ID flow - use existing implementation - return c.LoginWithBuilderID(ctx) - } - - // IDC flow - prompt for start URL and region - fmt.Println() - startURL := promptInput("? Enter Start URL", "") - if startURL == "" { - return nil, fmt.Errorf("start URL is required for IDC login") - } - - region := promptInput("? Enter Region", defaultIDCRegion) - - return c.LoginWithIDC(ctx, startURL, region) -} - -// RegisterClient registers a new OIDC client with AWS. -func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResponse, error) { - payload := map[string]interface{}{ - "clientName": "Kiro IDE", - "clientType": "public", - "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"}, - "grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"}, - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode) - } - - var result RegisterClientResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// StartDeviceAuthorization starts the device authorization flow. -func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID, clientSecret string) (*StartDeviceAuthResponse, error) { - payload := map[string]string{ - "clientId": clientID, - "clientSecret": clientSecret, - "startUrl": builderIDStartURL, - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/device_authorization", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode) - } - - var result StartDeviceAuthResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// CreateToken polls for the access token after user authorization. -func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret, deviceCode string) (*CreateTokenResponse, error) { - payload := map[string]string{ - "clientId": clientID, - "clientSecret": clientSecret, - "deviceCode": deviceCode, - "grantType": "urn:ietf:params:oauth:grant-type:device_code", - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - // Check for pending authorization - if resp.StatusCode == http.StatusBadRequest { - var errResp struct { - Error string `json:"error"` - } - if json.Unmarshal(respBody, &errResp) == nil { - if errResp.Error == "authorization_pending" { - return nil, ErrAuthorizationPending - } - if errResp.Error == "slow_down" { - return nil, ErrSlowDown - } - } - log.Debugf("create token failed: %s", string(respBody)) - return nil, fmt.Errorf("create token failed") - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode) - } - - var result CreateTokenResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// RefreshToken refreshes an access token using the refresh token. -func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret, refreshToken string) (*KiroTokenData, error) { - payload := map[string]string{ - "clientId": clientID, - "clientSecret": clientSecret, - "refreshToken": refreshToken, - "grantType": "refresh_token", - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode) - } - - var result CreateTokenResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second) - - return &KiroTokenData{ - AccessToken: result.AccessToken, - RefreshToken: result.RefreshToken, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: "builder-id", - Provider: "AWS", - ClientID: clientID, - ClientSecret: clientSecret, - }, nil -} - -// LoginWithBuilderID performs the full device code flow for AWS Builder ID. -func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) { - fmt.Println("\n╔══════════════════════════════════════════════════════════╗") - fmt.Println("║ Kiro Authentication (AWS Builder ID) ║") - fmt.Println("╚══════════════════════════════════════════════════════════╝") - - // Step 1: Register client - fmt.Println("\nRegistering client...") - regResp, err := c.RegisterClient(ctx) - if err != nil { - return nil, fmt.Errorf("failed to register client: %w", err) - } - log.Debugf("Client registered: %s", regResp.ClientID) - - // Step 2: Start device authorization - fmt.Println("Starting device authorization...") - authResp, err := c.StartDeviceAuthorization(ctx, regResp.ClientID, regResp.ClientSecret) - if err != nil { - return nil, fmt.Errorf("failed to start device auth: %w", err) - } - - // Step 3: Show user the verification URL - fmt.Printf("\n") - fmt.Println("════════════════════════════════════════════════════════════") - fmt.Printf(" Open this URL in your browser:\n") - fmt.Printf(" %s\n", authResp.VerificationURIComplete) - fmt.Println("════════════════════════════════════════════════════════════") - fmt.Printf("\n Or go to: %s\n", authResp.VerificationURI) - fmt.Printf(" And enter code: %s\n\n", authResp.UserCode) - - // Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito) - // Incognito mode enables multi-account support by bypassing cached sessions - if c.cfg != nil { - browser.SetIncognitoMode(c.cfg.IncognitoBrowser) - if !c.cfg.IncognitoBrowser { - log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.") - } else { - log.Debug("kiro: using incognito mode for multi-account support") - } - } else { - browser.SetIncognitoMode(true) // Default to incognito if no config - log.Debug("kiro: using incognito mode for multi-account support (default)") - } - - // Open browser using cross-platform browser package - if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil { - log.Warnf("Could not open browser automatically: %v", err) - fmt.Println(" Please open the URL manually in your browser.") - } else { - fmt.Println(" (Browser opened automatically)") - } - - // Step 4: Poll for token - fmt.Println("Waiting for authorization...") - - interval := pollInterval - if authResp.Interval > 0 { - interval = time.Duration(authResp.Interval) * time.Second - } - - deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second) - - for time.Now().Before(deadline) { - select { - case <-ctx.Done(): - browser.CloseBrowser() // Cleanup on cancel - return nil, ctx.Err() - case <-time.After(interval): - tokenResp, err := c.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode) - if err != nil { - if errors.Is(err, ErrAuthorizationPending) { - fmt.Print(".") - continue - } - if errors.Is(err, ErrSlowDown) { - interval += 5 * time.Second - continue - } - // Close browser on error before returning - browser.CloseBrowser() - return nil, fmt.Errorf("token creation failed: %w", err) - } - - fmt.Println("\n\n✓ Authorization successful!") - - // Close the browser window - if err := browser.CloseBrowser(); err != nil { - log.Debugf("Failed to close browser: %v", err) - } - - // Step 5: Get profile ARN from CodeWhisperer API - fmt.Println("Fetching profile information...") - profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken) - - // Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing) - email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken) - if email != "" { - fmt.Printf(" Logged in as: %s\n", email) - } - - expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) - - return &KiroTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - ProfileArn: profileArn, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: "builder-id", - Provider: "AWS", - ClientID: regResp.ClientID, - ClientSecret: regResp.ClientSecret, - Email: email, - }, nil - } - } - - // Close browser on timeout for better UX - if err := browser.CloseBrowser(); err != nil { - log.Debugf("Failed to close browser on timeout: %v", err) - } - return nil, fmt.Errorf("authorization timed out") -} - -// FetchUserEmail retrieves the user's email from AWS SSO OIDC userinfo endpoint. -// Falls back to JWT parsing if userinfo fails. -func (c *SSOOIDCClient) FetchUserEmail(ctx context.Context, accessToken string) string { - // Method 1: Try userinfo endpoint (standard OIDC) - email := c.tryUserInfoEndpoint(ctx, accessToken) - if email != "" { - return email - } - - // Method 2: Fallback to JWT parsing - return ExtractEmailFromJWT(accessToken) -} - -// tryUserInfoEndpoint attempts to get user info from AWS SSO OIDC userinfo endpoint. -func (c *SSOOIDCClient) tryUserInfoEndpoint(ctx context.Context, accessToken string) string { - req, err := http.NewRequestWithContext(ctx, http.MethodGet, ssoOIDCEndpoint+"/userinfo", nil) - if err != nil { - return "" - } - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - resp, err := c.httpClient.Do(req) - if err != nil { - log.Debugf("userinfo request failed: %v", err) - return "" - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - respBody, _ := io.ReadAll(resp.Body) - log.Debugf("userinfo endpoint returned status %d: %s", resp.StatusCode, string(respBody)) - return "" - } - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return "" - } - - log.Debugf("userinfo response: %s", string(respBody)) - - var userInfo struct { - Email string `json:"email"` - Sub string `json:"sub"` - PreferredUsername string `json:"preferred_username"` - Name string `json:"name"` - } - - if err := json.Unmarshal(respBody, &userInfo); err != nil { - return "" - } - - if userInfo.Email != "" { - return userInfo.Email - } - if userInfo.PreferredUsername != "" && strings.Contains(userInfo.PreferredUsername, "@") { - return userInfo.PreferredUsername - } - return "" -} - -// fetchProfileArn retrieves the profile ARN from CodeWhisperer API. -// This is needed for file naming since AWS SSO OIDC doesn't return profile info. -func (c *SSOOIDCClient) fetchProfileArn(ctx context.Context, accessToken string) string { - // Try ListProfiles API first - profileArn := c.tryListProfiles(ctx, accessToken) - if profileArn != "" { - return profileArn - } - - // Fallback: Try ListAvailableCustomizations - return c.tryListCustomizations(ctx, accessToken) -} - -func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string) string { - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - } - - body, err := json.Marshal(payload) - if err != nil { - return "" - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body))) - if err != nil { - return "" - } - - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListProfiles") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - resp, err := c.httpClient.Do(req) - if err != nil { - return "" - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != http.StatusOK { - log.Debugf("ListProfiles failed (status %d): %s", resp.StatusCode, string(respBody)) - return "" - } - - log.Debugf("ListProfiles response: %s", string(respBody)) - - var result struct { - Profiles []struct { - Arn string `json:"arn"` - } `json:"profiles"` - ProfileArn string `json:"profileArn"` - } - - if err := json.Unmarshal(respBody, &result); err != nil { - return "" - } - - if result.ProfileArn != "" { - return result.ProfileArn - } - - if len(result.Profiles) > 0 { - return result.Profiles[0].Arn - } - - return "" -} - -func (c *SSOOIDCClient) tryListCustomizations(ctx context.Context, accessToken string) string { - payload := map[string]interface{}{ - "origin": "AI_EDITOR", - } - - body, err := json.Marshal(payload) - if err != nil { - return "" - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body))) - if err != nil { - return "" - } - - req.Header.Set("Content-Type", "application/x-amz-json-1.0") - req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListAvailableCustomizations") - req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "application/json") - - resp, err := c.httpClient.Do(req) - if err != nil { - return "" - } - defer resp.Body.Close() - - respBody, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != http.StatusOK { - log.Debugf("ListAvailableCustomizations failed (status %d): %s", resp.StatusCode, string(respBody)) - return "" - } - - log.Debugf("ListAvailableCustomizations response: %s", string(respBody)) - - var result struct { - Customizations []struct { - Arn string `json:"arn"` - } `json:"customizations"` - ProfileArn string `json:"profileArn"` - } - - if err := json.Unmarshal(respBody, &result); err != nil { - return "" - } - - if result.ProfileArn != "" { - return result.ProfileArn - } - - if len(result.Customizations) > 0 { - return result.Customizations[0].Arn - } - - return "" -} - -// RegisterClientForAuthCode registers a new OIDC client for authorization code flow. -func (c *SSOOIDCClient) RegisterClientForAuthCode(ctx context.Context, redirectURI string) (*RegisterClientResponse, error) { - payload := map[string]interface{}{ - "clientName": "Kiro IDE", - "clientType": "public", - "scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"}, - "grantTypes": []string{"authorization_code", "refresh_token"}, - "redirectUris": []string{redirectURI}, - "issuerUrl": builderIDStartURL, - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("register client for auth code failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode) - } - - var result RegisterClientResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// AuthCodeCallbackResult contains the result from authorization code callback. -type AuthCodeCallbackResult struct { - Code string - State string - Error string -} - -// startAuthCodeCallbackServer starts a local HTTP server to receive the authorization code callback. -func (c *SSOOIDCClient) startAuthCodeCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthCodeCallbackResult, error) { - // Try to find an available port - listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", authCodeCallbackPort)) - if err != nil { - // Try with dynamic port - log.Warnf("sso oidc: default port %d is busy, falling back to dynamic port", authCodeCallbackPort) - listener, err = net.Listen("tcp", "127.0.0.1:0") - if err != nil { - return "", nil, fmt.Errorf("failed to start callback server: %w", err) - } - } - - port := listener.Addr().(*net.TCPAddr).Port - redirectURI := fmt.Sprintf("http://127.0.0.1:%d%s", port, authCodeCallbackPath) - resultChan := make(chan AuthCodeCallbackResult, 1) - - server := &http.Server{ - ReadHeaderTimeout: 10 * time.Second, - } - - mux := http.NewServeMux() - mux.HandleFunc(authCodeCallbackPath, func(w http.ResponseWriter, r *http.Request) { - code := r.URL.Query().Get("code") - state := r.URL.Query().Get("state") - errParam := r.URL.Query().Get("error") - - // Send response to browser - w.Header().Set("Content-Type", "text/html; charset=utf-8") - if errParam != "" { - w.WriteHeader(http.StatusBadRequest) - fmt.Fprintf(w, ` -Login Failed -

Login Failed

Error: %s

You can close this window.

`, html.EscapeString(errParam)) - resultChan <- AuthCodeCallbackResult{Error: errParam} - return - } - - if state != expectedState { - w.WriteHeader(http.StatusBadRequest) - fmt.Fprint(w, ` -Login Failed -

Login Failed

Invalid state parameter

You can close this window.

`) - resultChan <- AuthCodeCallbackResult{Error: "state mismatch"} - return - } - - fmt.Fprint(w, ` -Login Successful -

Login Successful!

You can close this window and return to the terminal.

-`) - resultChan <- AuthCodeCallbackResult{Code: code, State: state} - }) - - server.Handler = mux - - go func() { - if err := server.Serve(listener); err != nil && err != http.ErrServerClosed { - log.Debugf("auth code callback server error: %v", err) - } - }() - - go func() { - select { - case <-ctx.Done(): - case <-time.After(10 * time.Minute): - case <-resultChan: - } - _ = server.Shutdown(context.Background()) - }() - - return redirectURI, resultChan, nil -} - -// generatePKCEForAuthCode generates PKCE code verifier and challenge for authorization code flow. -func generatePKCEForAuthCode() (verifier, challenge string, err error) { - b := make([]byte, 32) - if _, err := rand.Read(b); err != nil { - return "", "", fmt.Errorf("failed to generate random bytes: %w", err) - } - verifier = base64.RawURLEncoding.EncodeToString(b) - h := sha256.Sum256([]byte(verifier)) - challenge = base64.RawURLEncoding.EncodeToString(h[:]) - return verifier, challenge, nil -} - -// generateStateForAuthCode generates a random state parameter. -func generateStateForAuthCode() (string, error) { - b := make([]byte, 16) - if _, err := rand.Read(b); err != nil { - return "", err - } - return base64.RawURLEncoding.EncodeToString(b), nil -} - -// CreateTokenWithAuthCode exchanges authorization code for tokens. -func (c *SSOOIDCClient) CreateTokenWithAuthCode(ctx context.Context, clientID, clientSecret, code, codeVerifier, redirectURI string) (*CreateTokenResponse, error) { - payload := map[string]string{ - "clientId": clientID, - "clientSecret": clientSecret, - "code": code, - "codeVerifier": codeVerifier, - "redirectUri": redirectURI, - "grantType": "authorization_code", - } - - body, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body))) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", kiroUserAgent) - - resp, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - if resp.StatusCode != http.StatusOK { - log.Debugf("create token with auth code failed (status %d): %s", resp.StatusCode, string(respBody)) - return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode) - } - - var result CreateTokenResponse - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, err - } - - return &result, nil -} - -// LoginWithBuilderIDAuthCode performs the authorization code flow for AWS Builder ID. -// This provides a better UX than device code flow as it uses automatic browser callback. -func (c *SSOOIDCClient) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTokenData, error) { - fmt.Println("\n╔══════════════════════════════════════════════════════════╗") - fmt.Println("║ Kiro Authentication (AWS Builder ID - Auth Code) ║") - fmt.Println("╚══════════════════════════════════════════════════════════╝") - - // Step 1: Generate PKCE and state - codeVerifier, codeChallenge, err := generatePKCEForAuthCode() - if err != nil { - return nil, fmt.Errorf("failed to generate PKCE: %w", err) - } - - state, err := generateStateForAuthCode() - if err != nil { - return nil, fmt.Errorf("failed to generate state: %w", err) - } - - // Step 2: Start callback server - fmt.Println("\nStarting callback server...") - redirectURI, resultChan, err := c.startAuthCodeCallbackServer(ctx, state) - if err != nil { - return nil, fmt.Errorf("failed to start callback server: %w", err) - } - log.Debugf("Callback server started, redirect URI: %s", redirectURI) - - // Step 3: Register client with auth code grant type - fmt.Println("Registering client...") - regResp, err := c.RegisterClientForAuthCode(ctx, redirectURI) - if err != nil { - return nil, fmt.Errorf("failed to register client: %w", err) - } - log.Debugf("Client registered: %s", regResp.ClientID) - - // Step 4: Build authorization URL - scopes := "codewhisperer:completions,codewhisperer:analysis,codewhisperer:conversations" - authURL := fmt.Sprintf("%s/authorize?response_type=code&client_id=%s&redirect_uri=%s&scopes=%s&state=%s&code_challenge=%s&code_challenge_method=S256", - ssoOIDCEndpoint, - regResp.ClientID, - redirectURI, - scopes, - state, - codeChallenge, - ) - - // Step 5: Open browser - fmt.Println("\n════════════════════════════════════════════════════════════") - fmt.Println(" Opening browser for authentication...") - fmt.Println("════════════════════════════════════════════════════════════") - fmt.Printf("\n URL: %s\n\n", authURL) - - // Set incognito mode - if c.cfg != nil { - browser.SetIncognitoMode(c.cfg.IncognitoBrowser) - } else { - browser.SetIncognitoMode(true) - } - - if err := browser.OpenURL(authURL); err != nil { - log.Warnf("Could not open browser automatically: %v", err) - fmt.Println(" ⚠ Could not open browser automatically.") - fmt.Println(" Please open the URL above in your browser manually.") - } else { - fmt.Println(" (Browser opened automatically)") - } - - fmt.Println("\n Waiting for authorization callback...") - - // Step 6: Wait for callback - select { - case <-ctx.Done(): - browser.CloseBrowser() - return nil, ctx.Err() - case <-time.After(10 * time.Minute): - browser.CloseBrowser() - return nil, fmt.Errorf("authorization timed out") - case result := <-resultChan: - if result.Error != "" { - browser.CloseBrowser() - return nil, fmt.Errorf("authorization failed: %s", result.Error) - } - - fmt.Println("\n✓ Authorization received!") - - // Close browser - if err := browser.CloseBrowser(); err != nil { - log.Debugf("Failed to close browser: %v", err) - } - - // Step 7: Exchange code for tokens - fmt.Println("Exchanging code for tokens...") - tokenResp, err := c.CreateTokenWithAuthCode(ctx, regResp.ClientID, regResp.ClientSecret, result.Code, codeVerifier, redirectURI) - if err != nil { - return nil, fmt.Errorf("failed to exchange code for tokens: %w", err) - } - - fmt.Println("\n✓ Authentication successful!") - - // Step 8: Get profile ARN - fmt.Println("Fetching profile information...") - profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken) - - // Fetch user email (tries CodeWhisperer API first, then userinfo endpoint, then JWT parsing) - email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken) - if email != "" { - fmt.Printf(" Logged in as: %s\n", email) - } - - expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) - - return &KiroTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - ProfileArn: profileArn, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: "builder-id", - Provider: "AWS", - ClientID: regResp.ClientID, - ClientSecret: regResp.ClientSecret, - Email: email, - }, nil - } -} diff --git a/internal/translator/kiro/common/utf8_stream.go b/internal/translator/kiro/common/utf8_stream.go deleted file mode 100644 index b8d24c82..00000000 --- a/internal/translator/kiro/common/utf8_stream.go +++ /dev/null @@ -1,97 +0,0 @@ -package common - -import ( - "unicode/utf8" -) - -type UTF8StreamParser struct { - buffer []byte -} - -func NewUTF8StreamParser() *UTF8StreamParser { - return &UTF8StreamParser{ - buffer: make([]byte, 0, 64), - } -} - -func (p *UTF8StreamParser) Write(data []byte) { - p.buffer = append(p.buffer, data...) -} - -func (p *UTF8StreamParser) Read() (string, bool) { - if len(p.buffer) == 0 { - return "", false - } - - validLen := p.findValidUTF8End(p.buffer) - if validLen == 0 { - return "", false - } - - result := string(p.buffer[:validLen]) - p.buffer = p.buffer[validLen:] - - return result, true -} - -func (p *UTF8StreamParser) Flush() string { - if len(p.buffer) == 0 { - return "" - } - result := string(p.buffer) - p.buffer = p.buffer[:0] - return result -} - -func (p *UTF8StreamParser) Reset() { - p.buffer = p.buffer[:0] -} - -func (p *UTF8StreamParser) findValidUTF8End(data []byte) int { - if len(data) == 0 { - return 0 - } - - end := len(data) - for i := 1; i <= 3 && i <= len(data); i++ { - b := data[len(data)-i] - if b&0x80 == 0 { - break - } - if b&0xC0 == 0xC0 { - size := p.utf8CharSize(b) - available := i - if size > available { - end = len(data) - i - } - break - } - } - - if end > 0 && !utf8.Valid(data[:end]) { - for i := end - 1; i >= 0; i-- { - if utf8.Valid(data[:i+1]) { - return i + 1 - } - } - return 0 - } - - return end -} - -func (p *UTF8StreamParser) utf8CharSize(b byte) int { - if b&0x80 == 0 { - return 1 - } - if b&0xE0 == 0xC0 { - return 2 - } - if b&0xF0 == 0xE0 { - return 3 - } - if b&0xF8 == 0xF0 { - return 4 - } - return 1 -} diff --git a/internal/translator/kiro/common/utf8_stream_test.go b/internal/translator/kiro/common/utf8_stream_test.go deleted file mode 100644 index 23e80989..00000000 --- a/internal/translator/kiro/common/utf8_stream_test.go +++ /dev/null @@ -1,402 +0,0 @@ -package common - -import ( - "strings" - "sync" - "testing" - "unicode/utf8" -) - -func TestNewUTF8StreamParser(t *testing.T) { - p := NewUTF8StreamParser() - if p == nil { - t.Fatal("expected non-nil UTF8StreamParser") - } - if p.buffer == nil { - t.Error("expected non-nil buffer") - } -} - -func TestWrite(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("hello")) - - result, ok := p.Read() - if !ok { - t.Error("expected ok to be true") - } - if result != "hello" { - t.Errorf("expected 'hello', got '%s'", result) - } -} - -func TestWrite_MultipleWrites(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("hel")) - p.Write([]byte("lo")) - - result, ok := p.Read() - if !ok { - t.Error("expected ok to be true") - } - if result != "hello" { - t.Errorf("expected 'hello', got '%s'", result) - } -} - -func TestRead_EmptyBuffer(t *testing.T) { - p := NewUTF8StreamParser() - result, ok := p.Read() - if ok { - t.Error("expected ok to be false for empty buffer") - } - if result != "" { - t.Errorf("expected empty string, got '%s'", result) - } -} - -func TestRead_IncompleteUTF8(t *testing.T) { - p := NewUTF8StreamParser() - - // Write incomplete multi-byte UTF-8 character - // 中 (U+4E2D) = E4 B8 AD - p.Write([]byte{0xE4, 0xB8}) - - result, ok := p.Read() - if ok { - t.Error("expected ok to be false for incomplete UTF-8") - } - if result != "" { - t.Errorf("expected empty string, got '%s'", result) - } - - // Complete the character - p.Write([]byte{0xAD}) - result, ok = p.Read() - if !ok { - t.Error("expected ok to be true after completing UTF-8") - } - if result != "中" { - t.Errorf("expected '中', got '%s'", result) - } -} - -func TestRead_MixedASCIIAndUTF8(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("Hello 世界")) - - result, ok := p.Read() - if !ok { - t.Error("expected ok to be true") - } - if result != "Hello 世界" { - t.Errorf("expected 'Hello 世界', got '%s'", result) - } -} - -func TestRead_PartialMultibyteAtEnd(t *testing.T) { - p := NewUTF8StreamParser() - // "Hello" + partial "世" (E4 B8 96) - p.Write([]byte("Hello")) - p.Write([]byte{0xE4, 0xB8}) - - result, ok := p.Read() - if !ok { - t.Error("expected ok to be true for valid portion") - } - if result != "Hello" { - t.Errorf("expected 'Hello', got '%s'", result) - } - - // Complete the character - p.Write([]byte{0x96}) - result, ok = p.Read() - if !ok { - t.Error("expected ok to be true after completing") - } - if result != "世" { - t.Errorf("expected '世', got '%s'", result) - } -} - -func TestFlush(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("hello")) - - result := p.Flush() - if result != "hello" { - t.Errorf("expected 'hello', got '%s'", result) - } - - // Verify buffer is cleared - result2, ok := p.Read() - if ok { - t.Error("expected ok to be false after flush") - } - if result2 != "" { - t.Errorf("expected empty string after flush, got '%s'", result2) - } -} - -func TestFlush_EmptyBuffer(t *testing.T) { - p := NewUTF8StreamParser() - result := p.Flush() - if result != "" { - t.Errorf("expected empty string, got '%s'", result) - } -} - -func TestFlush_IncompleteUTF8(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte{0xE4, 0xB8}) - - result := p.Flush() - // Flush returns everything including incomplete bytes - if len(result) != 2 { - t.Errorf("expected 2 bytes flushed, got %d", len(result)) - } -} - -func TestReset(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("hello")) - p.Reset() - - result, ok := p.Read() - if ok { - t.Error("expected ok to be false after reset") - } - if result != "" { - t.Errorf("expected empty string after reset, got '%s'", result) - } -} - -func TestUtf8CharSize(t *testing.T) { - p := NewUTF8StreamParser() - - testCases := []struct { - b byte - expected int - }{ - {0x00, 1}, // ASCII - {0x7F, 1}, // ASCII max - {0xC0, 2}, // 2-byte start - {0xDF, 2}, // 2-byte start - {0xE0, 3}, // 3-byte start - {0xEF, 3}, // 3-byte start - {0xF0, 4}, // 4-byte start - {0xF7, 4}, // 4-byte start - {0x80, 1}, // Continuation byte (fallback) - } - - for _, tc := range testCases { - size := p.utf8CharSize(tc.b) - if size != tc.expected { - t.Errorf("utf8CharSize(0x%X) = %d, expected %d", tc.b, size, tc.expected) - } - } -} - -func TestStreamingScenario(t *testing.T) { - p := NewUTF8StreamParser() - - // Simulate streaming: "Hello, 世界! 🌍" - chunks := [][]byte{ - []byte("Hello, "), - {0xE4, 0xB8}, // partial 世 - {0x96, 0xE7}, // complete 世, partial 界 - {0x95, 0x8C}, // complete 界 - []byte("! "), - {0xF0, 0x9F}, // partial 🌍 - {0x8C, 0x8D}, // complete 🌍 - } - - var results []string - for _, chunk := range chunks { - p.Write(chunk) - if result, ok := p.Read(); ok { - results = append(results, result) - } - } - - combined := strings.Join(results, "") - if combined != "Hello, 世界! 🌍" { - t.Errorf("expected 'Hello, 世界! 🌍', got '%s'", combined) - } -} - -func TestValidUTF8Output(t *testing.T) { - p := NewUTF8StreamParser() - - testStrings := []string{ - "Hello World", - "你好世界", - "こんにちは", - "🎉🎊🎁", - "Mixed 混合 Текст ტექსტი", - } - - for _, s := range testStrings { - p.Reset() - p.Write([]byte(s)) - result, ok := p.Read() - if !ok { - t.Errorf("expected ok for '%s'", s) - } - if !utf8.ValidString(result) { - t.Errorf("invalid UTF-8 output for input '%s'", s) - } - if result != s { - t.Errorf("expected '%s', got '%s'", s, result) - } - } -} - -func TestLargeData(t *testing.T) { - p := NewUTF8StreamParser() - - // Generate large UTF-8 string - var builder strings.Builder - for i := 0; i < 1000; i++ { - builder.WriteString("Hello 世界! ") - } - largeString := builder.String() - - p.Write([]byte(largeString)) - result, ok := p.Read() - if !ok { - t.Error("expected ok for large data") - } - if result != largeString { - t.Error("large data mismatch") - } -} - -func TestByteByByteWriting(t *testing.T) { - p := NewUTF8StreamParser() - input := "Hello 世界" - inputBytes := []byte(input) - - var results []string - for _, b := range inputBytes { - p.Write([]byte{b}) - if result, ok := p.Read(); ok { - results = append(results, result) - } - } - - combined := strings.Join(results, "") - if combined != input { - t.Errorf("expected '%s', got '%s'", input, combined) - } -} - -func TestEmoji4ByteUTF8(t *testing.T) { - p := NewUTF8StreamParser() - - // 🎉 = F0 9F 8E 89 - emoji := "🎉" - emojiBytes := []byte(emoji) - - for i := 0; i < len(emojiBytes)-1; i++ { - p.Write(emojiBytes[i : i+1]) - result, ok := p.Read() - if ok && result != "" { - t.Errorf("unexpected output before emoji complete: '%s'", result) - } - } - - p.Write(emojiBytes[len(emojiBytes)-1:]) - result, ok := p.Read() - if !ok { - t.Error("expected ok after completing emoji") - } - if result != emoji { - t.Errorf("expected '%s', got '%s'", emoji, result) - } -} - -func TestContinuationBytesOnly(t *testing.T) { - p := NewUTF8StreamParser() - - // Write only continuation bytes (invalid UTF-8) - p.Write([]byte{0x80, 0x80, 0x80}) - - result, ok := p.Read() - // Should handle gracefully - either return nothing or return the bytes - _ = result - _ = ok -} - -func TestUTF8StreamParser_ConcurrentSafety(t *testing.T) { - // Note: UTF8StreamParser doesn't have built-in locks, - // so this test verifies it works with external synchronization - p := NewUTF8StreamParser() - var mu sync.Mutex - const numGoroutines = 10 - const numOperations = 100 - - var wg sync.WaitGroup - wg.Add(numGoroutines) - - for i := 0; i < numGoroutines; i++ { - go func() { - defer wg.Done() - for j := 0; j < numOperations; j++ { - mu.Lock() - switch j % 4 { - case 0: - p.Write([]byte("test")) - case 1: - p.Read() - case 2: - p.Flush() - case 3: - p.Reset() - } - mu.Unlock() - } - }() - } - - wg.Wait() -} - -func TestConsecutiveReads(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("hello")) - - result1, ok1 := p.Read() - if !ok1 || result1 != "hello" { - t.Error("first read failed") - } - - result2, ok2 := p.Read() - if ok2 || result2 != "" { - t.Error("second read should return empty") - } -} - -func TestFlushThenWrite(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("first")) - p.Flush() - p.Write([]byte("second")) - - result, ok := p.Read() - if !ok || result != "second" { - t.Errorf("expected 'second', got '%s'", result) - } -} - -func TestResetThenWrite(t *testing.T) { - p := NewUTF8StreamParser() - p.Write([]byte("first")) - p.Reset() - p.Write([]byte("second")) - - result, ok := p.Read() - if !ok || result != "second" { - t.Errorf("expected 'second', got '%s'", result) - } -} diff --git a/sdk/auth/kiro.go.bak b/sdk/auth/kiro.go.bak deleted file mode 100644 index b75cd28e..00000000 --- a/sdk/auth/kiro.go.bak +++ /dev/null @@ -1,470 +0,0 @@ -package auth - -import ( - "context" - "fmt" - "strings" - "time" - - kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro" - "github.com/router-for-me/CLIProxyAPI/v6/internal/config" - coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" -) - -// extractKiroIdentifier extracts a meaningful identifier for file naming. -// Returns account name if provided, otherwise profile ARN ID. -// All extracted values are sanitized to prevent path injection attacks. -func extractKiroIdentifier(accountName, profileArn string) string { - // Priority 1: Use account name if provided - if accountName != "" { - return kiroauth.SanitizeEmailForFilename(accountName) - } - - // Priority 2: Use profile ARN ID part (sanitized to prevent path injection) - if profileArn != "" { - parts := strings.Split(profileArn, "/") - if len(parts) >= 2 { - // Sanitize the ARN component to prevent path traversal - return kiroauth.SanitizeEmailForFilename(parts[len(parts)-1]) - } - } - - // Fallback: timestamp - return fmt.Sprintf("%d", time.Now().UnixNano()%100000) -} - -// KiroAuthenticator implements OAuth authentication for Kiro with Google login. -type KiroAuthenticator struct{} - -// NewKiroAuthenticator constructs a Kiro authenticator. -func NewKiroAuthenticator() *KiroAuthenticator { - return &KiroAuthenticator{} -} - -// Provider returns the provider key for the authenticator. -func (a *KiroAuthenticator) Provider() string { - return "kiro" -} - -// RefreshLead indicates how soon before expiry a refresh should be attempted. -// Set to 5 minutes to match Antigravity and avoid frequent refresh checks while still ensuring timely token refresh. -func (a *KiroAuthenticator) RefreshLead() *time.Duration { - d := 5 * time.Minute - return &d -} - -// createAuthRecord creates an auth record from token data. -func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, source string) (*coreauth.Auth, error) { - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) - - // Determine label based on auth method - label := fmt.Sprintf("kiro-%s", source) - if tokenData.AuthMethod == "idc" { - label = "kiro-idc" - } - - now := time.Now() - fileName := fmt.Sprintf("%s-%s.json", label, idPart) - - metadata := map[string]any{ - "type": "kiro", - "access_token": tokenData.AccessToken, - "refresh_token": tokenData.RefreshToken, - "profile_arn": tokenData.ProfileArn, - "expires_at": tokenData.ExpiresAt, - "auth_method": tokenData.AuthMethod, - "provider": tokenData.Provider, - "client_id": tokenData.ClientID, - "client_secret": tokenData.ClientSecret, - "email": tokenData.Email, - } - - // Add IDC-specific fields if present - if tokenData.StartURL != "" { - metadata["start_url"] = tokenData.StartURL - } - if tokenData.Region != "" { - metadata["region"] = tokenData.Region - } - - attributes := map[string]string{ - "profile_arn": tokenData.ProfileArn, - "source": source, - "email": tokenData.Email, - } - - // Add IDC-specific attributes if present - if tokenData.AuthMethod == "idc" { - attributes["source"] = "aws-idc" - if tokenData.StartURL != "" { - attributes["start_url"] = tokenData.StartURL - } - if tokenData.Region != "" { - attributes["region"] = tokenData.Region - } - } - - record := &coreauth.Auth{ - ID: fileName, - Provider: "kiro", - FileName: fileName, - Label: label, - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Metadata: metadata, - Attributes: attributes, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), - } - - if tokenData.Email != "" { - fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email) - } else { - fmt.Println("\n✓ Kiro authentication completed successfully!") - } - - return record, nil -} - -// Login performs OAuth login for Kiro with AWS (Builder ID or IDC). -// This shows a method selection prompt and handles both flows. -func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { - if cfg == nil { - return nil, fmt.Errorf("kiro auth: configuration is required") - } - - // Use the unified method selection flow (Builder ID or IDC) - ssoClient := kiroauth.NewSSOOIDCClient(cfg) - tokenData, err := ssoClient.LoginWithMethodSelection(ctx) - if err != nil { - return nil, fmt.Errorf("login failed: %w", err) - } - - return a.createAuthRecord(tokenData, "aws") -} - -// LoginWithAuthCode performs OAuth login for Kiro with AWS Builder ID using authorization code flow. -// This provides a better UX than device code flow as it uses automatic browser callback. -func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { - if cfg == nil { - return nil, fmt.Errorf("kiro auth: configuration is required") - } - - oauth := kiroauth.NewKiroOAuth(cfg) - - // Use AWS Builder ID authorization code flow - tokenData, err := oauth.LoginWithBuilderIDAuthCode(ctx) - if err != nil { - return nil, fmt.Errorf("login failed: %w", err) - } - - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) - - now := time.Now() - fileName := fmt.Sprintf("kiro-aws-%s.json", idPart) - - record := &coreauth.Auth{ - ID: fileName, - Provider: "kiro", - FileName: fileName, - Label: "kiro-aws", - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Metadata: map[string]any{ - "type": "kiro", - "access_token": tokenData.AccessToken, - "refresh_token": tokenData.RefreshToken, - "profile_arn": tokenData.ProfileArn, - "expires_at": tokenData.ExpiresAt, - "auth_method": tokenData.AuthMethod, - "provider": tokenData.Provider, - "client_id": tokenData.ClientID, - "client_secret": tokenData.ClientSecret, - "email": tokenData.Email, - }, - Attributes: map[string]string{ - "profile_arn": tokenData.ProfileArn, - "source": "aws-builder-id-authcode", - "email": tokenData.Email, - }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), - } - - if tokenData.Email != "" { - fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email) - } else { - fmt.Println("\n✓ Kiro authentication completed successfully!") - } - - return record, nil -} - -// LoginWithGoogle performs OAuth login for Kiro with Google. -// This uses a custom protocol handler (kiro://) to receive the callback. -func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { - if cfg == nil { - return nil, fmt.Errorf("kiro auth: configuration is required") - } - - oauth := kiroauth.NewKiroOAuth(cfg) - - // Use Google OAuth flow with protocol handler - tokenData, err := oauth.LoginWithGoogle(ctx) - if err != nil { - return nil, fmt.Errorf("google login failed: %w", err) - } - - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) - - now := time.Now() - fileName := fmt.Sprintf("kiro-google-%s.json", idPart) - - record := &coreauth.Auth{ - ID: fileName, - Provider: "kiro", - FileName: fileName, - Label: "kiro-google", - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Metadata: map[string]any{ - "type": "kiro", - "access_token": tokenData.AccessToken, - "refresh_token": tokenData.RefreshToken, - "profile_arn": tokenData.ProfileArn, - "expires_at": tokenData.ExpiresAt, - "auth_method": tokenData.AuthMethod, - "provider": tokenData.Provider, - "email": tokenData.Email, - }, - Attributes: map[string]string{ - "profile_arn": tokenData.ProfileArn, - "source": "google-oauth", - "email": tokenData.Email, - }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), - } - - if tokenData.Email != "" { - fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email) - } else { - fmt.Println("\n✓ Kiro Google authentication completed successfully!") - } - - return record, nil -} - -// LoginWithGitHub performs OAuth login for Kiro with GitHub. -// This uses a custom protocol handler (kiro://) to receive the callback. -func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { - if cfg == nil { - return nil, fmt.Errorf("kiro auth: configuration is required") - } - - oauth := kiroauth.NewKiroOAuth(cfg) - - // Use GitHub OAuth flow with protocol handler - tokenData, err := oauth.LoginWithGitHub(ctx) - if err != nil { - return nil, fmt.Errorf("github login failed: %w", err) - } - - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) - - now := time.Now() - fileName := fmt.Sprintf("kiro-github-%s.json", idPart) - - record := &coreauth.Auth{ - ID: fileName, - Provider: "kiro", - FileName: fileName, - Label: "kiro-github", - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Metadata: map[string]any{ - "type": "kiro", - "access_token": tokenData.AccessToken, - "refresh_token": tokenData.RefreshToken, - "profile_arn": tokenData.ProfileArn, - "expires_at": tokenData.ExpiresAt, - "auth_method": tokenData.AuthMethod, - "provider": tokenData.Provider, - "email": tokenData.Email, - }, - Attributes: map[string]string{ - "profile_arn": tokenData.ProfileArn, - "source": "github-oauth", - "email": tokenData.Email, - }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), - } - - if tokenData.Email != "" { - fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email) - } else { - fmt.Println("\n✓ Kiro GitHub authentication completed successfully!") - } - - return record, nil -} - -// ImportFromKiroIDE imports token from Kiro IDE's token file. -func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.Config) (*coreauth.Auth, error) { - tokenData, err := kiroauth.LoadKiroIDEToken() - if err != nil { - return nil, fmt.Errorf("failed to load Kiro IDE token: %w", err) - } - - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Extract email from JWT if not already set (for imported tokens) - if tokenData.Email == "" { - tokenData.Email = kiroauth.ExtractEmailFromJWT(tokenData.AccessToken) - } - - // Extract identifier for file naming - idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn) - // Sanitize provider to prevent path traversal (defense-in-depth) - provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider))) - if provider == "" { - provider = "imported" // Fallback for legacy tokens without provider - } - - now := time.Now() - fileName := fmt.Sprintf("kiro-%s-%s.json", provider, idPart) - - record := &coreauth.Auth{ - ID: fileName, - Provider: "kiro", - FileName: fileName, - Label: fmt.Sprintf("kiro-%s", provider), - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Metadata: map[string]any{ - "type": "kiro", - "access_token": tokenData.AccessToken, - "refresh_token": tokenData.RefreshToken, - "profile_arn": tokenData.ProfileArn, - "expires_at": tokenData.ExpiresAt, - "auth_method": tokenData.AuthMethod, - "provider": tokenData.Provider, - "email": tokenData.Email, - }, - Attributes: map[string]string{ - "profile_arn": tokenData.ProfileArn, - "source": "kiro-ide-import", - "email": tokenData.Email, - }, - // NextRefreshAfter is aligned with RefreshLead (5min) - NextRefreshAfter: expiresAt.Add(-5 * time.Minute), - } - - // Display the email if extracted - if tokenData.Email != "" { - fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s, Account: %s)\n", tokenData.Provider, tokenData.Email) - } else { - fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s)\n", tokenData.Provider) - } - - return record, nil -} - -// Refresh refreshes an expired Kiro token using AWS SSO OIDC. -func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, auth *coreauth.Auth) (*coreauth.Auth, error) { - if auth == nil || auth.Metadata == nil { - return nil, fmt.Errorf("invalid auth record") - } - - refreshToken, ok := auth.Metadata["refresh_token"].(string) - if !ok || refreshToken == "" { - return nil, fmt.Errorf("refresh token not found") - } - - clientID, _ := auth.Metadata["client_id"].(string) - clientSecret, _ := auth.Metadata["client_secret"].(string) - authMethod, _ := auth.Metadata["auth_method"].(string) - startURL, _ := auth.Metadata["start_url"].(string) - region, _ := auth.Metadata["region"].(string) - - var tokenData *kiroauth.KiroTokenData - var err error - - ssoClient := kiroauth.NewSSOOIDCClient(cfg) - - // Use SSO OIDC refresh for AWS Builder ID or IDC, otherwise use Kiro's OAuth refresh endpoint - switch { - case clientID != "" && clientSecret != "" && authMethod == "idc" && region != "": - // IDC refresh with region-specific endpoint - tokenData, err = ssoClient.RefreshTokenWithRegion(ctx, clientID, clientSecret, refreshToken, region, startURL) - case clientID != "" && clientSecret != "" && authMethod == "builder-id": - // Builder ID refresh with default endpoint - tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken) - default: - // Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub) - oauth := kiroauth.NewKiroOAuth(cfg) - tokenData, err = oauth.RefreshToken(ctx, refreshToken) - } - - if err != nil { - return nil, fmt.Errorf("token refresh failed: %w", err) - } - - // Parse expires_at - expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt) - if err != nil { - expiresAt = time.Now().Add(1 * time.Hour) - } - - // Clone auth to avoid mutating the input parameter - updated := auth.Clone() - now := time.Now() - updated.UpdatedAt = now - updated.LastRefreshedAt = now - updated.Metadata["access_token"] = tokenData.AccessToken - updated.Metadata["refresh_token"] = tokenData.RefreshToken - updated.Metadata["expires_at"] = tokenData.ExpiresAt - updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization - // NextRefreshAfter is aligned with RefreshLead (5min) - updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute) - - return updated, nil -}