package executor import (	"bufio"	"bytes"	"context"	"encoding/base64"	"encoding/binary"	"encoding/json"	"errors"	"fmt"	"io"	"net"	"net/http"	"os"	"path/filepath"	"strings"	"sync"	"syscall"	"time" 	"github.com/google/uuid"	kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"	kiroclaude "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/claude"	kirocommon "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/common"	kiroopenai "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/openai"	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"	"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"	log "github.com/sirupsen/logrus") const (	// Kiro API common constants	kiroContentType  = "application/json"	kiroAcceptStream = "*/*" 	// Event Stream frame size constants for boundary protection	// AWS Event Stream binary format: prelude (12 bytes) + headers + payload + message_crc (4 bytes)	// Prelude consists of: total_length (4) + headers_length (4) + prelude_crc (4)	minEventStreamFrameSize = 16       // Minimum: 4(total_len) + 4(headers_len) + 4(prelude_crc) + 4(message_crc)	maxEventStreamMsgSize   = 10 << 20 // Maximum message length: 10MB 	// Event Stream error type constants	ErrStreamFatal     = "fatal"     // Connection/authentication errors, not recoverable	ErrStreamMalformed = "malformed" // Format errors, data cannot be parsed 	// kiroUserAgent matches Amazon Q CLI style for User-Agent header	kiroUserAgent = "aws-sdk-rust/1.3.9 os/macos lang/rust/1.87.0"	// kiroFullUserAgent is the complete x-amz-user-agent header (Amazon Q CLI style)	kiroFullUserAgent = "aws-sdk-rust/1.3.9 ua/2.1 api/ssooidc/1.88.0 os/macos lang/rust/1.87.0 m/E app/AmazonQ-For-CLI" 	// Kiro IDE style headers for IDC auth	kiroIDEUserAgent     = "aws-sdk-js/1.0.27 ua/2.1 os/win32#10.0.19044 lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E"	kiroIDEAmzUserAgent  = "aws-sdk-js/1.0.27"	kiroIDEAgentModeVibe = "vibe" 	// Socket retry configuration constants	// 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 streamingvar (	usageUpdateCharThreshold = 5000             // Send usage update every 5000 characters	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 requestsvar (	globalFingerprintManager     *kiroauth.FingerprintManager	globalFingerprintManagerOnce sync.Once) // getGlobalFingerprintManager returns the global FingerprintManager instancefunc 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 connectionsfunc 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.//// Based on reference implementations:// - amq2api-main: Uses Amazon Q endpoint with CLI origin and AmazonQDeveloperStreamingService target// - AIClient-2-API: Uses CodeWhisperer endpoint with AI_EDITOR origin and AmazonCodeWhispererStreamingService targettype kiroEndpointConfig struct {	URL       string // Endpoint URL	Origin    string // Request Origin: "CLI" for Amazon Q quota, "AI_EDITOR" for Kiro IDE quota	AmzTarget string // X-Amz-Target header value	Name      string // Endpoint name for logging} // kiroDefaultRegion is the default AWS region for Kiro API endpoints.// Used when no region is specified in auth metadata.const kiroDefaultRegion = "us-east-1" // extractRegionFromProfileARN extracts the AWS region from a ProfileARN.// ARN format: arn:aws:codewhisperer:REGION:ACCOUNT:profile/PROFILE_ID// Returns empty string if region cannot be extracted.func extractRegionFromProfileARN(profileArn string) string {	if profileArn == "" {		return ""	}	parts := strings.Split(profileArn, ":")	if len(parts) >= 4 && parts[3] != "" {		return parts[3]	}	return ""} // buildKiroEndpointConfigs creates endpoint configurations for the specified region.// This enables dynamic region support for Enterprise/IdC users in non-us-east-1 regions.//// Uses Q endpoint (q.{region}.amazonaws.com) as primary for ALL auth types:// - Works universally across all AWS regions (CodeWhisperer endpoint only exists in us-east-1)// - Uses /generateAssistantResponse path with AI_EDITOR origin// - Does NOT require X-Amz-Target header//// The AmzTarget field is kept for backward compatibility but should be empty// to indicate that the header should NOT be set.func buildKiroEndpointConfigs(region string) []kiroEndpointConfig {	if region == "" {		region = kiroDefaultRegion	}	return []kiroEndpointConfig{		{			// Primary: Q endpoint - works for all regions and auth types			URL:       fmt.Sprintf("https://q.%s.amazonaws.com/generateAssistantResponse", region),			Origin:    "AI_EDITOR",			AmzTarget: "", // Empty = don't set X-Amz-Target header			Name:      "AmazonQ",		},		{			// Fallback: CodeWhisperer endpoint (legacy, only works in us-east-1)			URL:       fmt.Sprintf("https://codewhisperer.%s.amazonaws.com/generateAssistantResponse", region),			Origin:    "AI_EDITOR",			AmzTarget: "AmazonCodeWhispererStreamingService.GenerateAssistantResponse",			Name:      "CodeWhisperer",		},	}} // kiroEndpointConfigs is kept for backward compatibility with default us-east-1 region.// Prefer using buildKiroEndpointConfigs(region) for dynamic region support.var kiroEndpointConfigs = buildKiroEndpointConfigs(kiroDefaultRegion) // getKiroEndpointConfigs returns the list of Kiro API endpoint configurations to try in order.// Supports dynamic region based on auth metadata "api_region", "profile_arn", or "region" field.// Supports reordering based on "preferred_endpoint" in auth metadata/attributes.//// Region priority:// 1. auth.Metadata["api_region"] - explicit API region override// 2. ProfileARN region - extracted from arn:aws:service:REGION:account:resource// 3. kiroDefaultRegion (us-east-1) - fallback// Note: OIDC "region" is NOT used - it's for token refresh, not API callsfunc getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig {	if auth == nil {		return kiroEndpointConfigs	} 	// Determine API region with priority: api_region > profile_arn > region > default	region := kiroDefaultRegion	regionSource := "default" 	if auth.Metadata != nil {		// Priority 1: Explicit api_region override		if r, ok := auth.Metadata["api_region"].(string); ok && r != "" {			region = r			regionSource = "api_region"		} else {			// Priority 2: Extract from ProfileARN			if profileArn, ok := auth.Metadata["profile_arn"].(string); ok && profileArn != "" {				if arnRegion := extractRegionFromProfileARN(profileArn); arnRegion != "" {					region = arnRegion					regionSource = "profile_arn"				}			}			// Note: OIDC "region" field is NOT used for API endpoint			// Kiro API only exists in us-east-1, while OIDC region can vary (e.g., ap-northeast-2)			// Using OIDC region for API calls causes DNS failures		}	} 	log.Debugf("kiro: using region %s (source: %s)", region, regionSource) 	// Build endpoint configs for the specified region	endpointConfigs := buildKiroEndpointConfigs(region) 	// For IDC auth, use Q endpoint with AI_EDITOR origin	// IDC tokens work with Q endpoint using Bearer auth	// The difference is only in how tokens are refreshed (OIDC with clientId/clientSecret for IDC)	// NOT in how API calls are made - both Social and IDC use the same endpoint/origin	if auth.Metadata != nil {		authMethod, _ := auth.Metadata["auth_method"].(string)		if strings.ToLower(authMethod) == "idc" {			log.Debugf("kiro: IDC auth, using Q endpoint (region: %s)", region)			return endpointConfigs		}	} 	// Check for preference	var preference string	if auth.Metadata != nil {		if p, ok := auth.Metadata["preferred_endpoint"].(string); ok {			preference = p		}	}	// Check attributes as fallback (e.g. from HTTP headers)	if preference == "" && auth.Attributes != nil {		preference = auth.Attributes["preferred_endpoint"]	} 	if preference == "" {		return endpointConfigs	} 	preference = strings.ToLower(strings.TrimSpace(preference)) 	// Create new slice to avoid modifying global state	var sorted []kiroEndpointConfig	var remaining []kiroEndpointConfig 	for _, cfg := range endpointConfigs {		name := strings.ToLower(cfg.Name)		// Check for matches		// CodeWhisperer aliases: codewhisperer, ide		// AmazonQ aliases: amazonq, q, cli		isMatch := false		if (preference == "codewhisperer" || preference == "ide") && name == "codewhisperer" {			isMatch = true		} else if (preference == "amazonq" || preference == "q" || preference == "cli") && name == "amazonq" {			isMatch = true		} 		if isMatch {			sorted = append(sorted, cfg)		} else {			remaining = append(remaining, cfg)		}	} 	// If preference didn't match anything, return default	if len(sorted) == 0 {		return endpointConfigs	} 	// Combine: preferred first, then others	return append(sorted, remaining...)} // KiroExecutor handles requests to AWS CodeWhisperer (Kiro) API.type KiroExecutor struct {	cfg       *config.Config	refreshMu sync.Mutex // Serializes token refresh operations to prevent race conditions}

Lines 1–500 of 4101. Load more