package executor import (	"bufio"	"bytes"	"context"	"encoding/base64"	"encoding/binary"	"encoding/json"	"errors"	"fmt"	"io"	"net"	"net/http"	"os"	"path/filepath"	"strings"	"sync"	"sync/atomic"	"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 	// kiroIDEAgentMode is the agent mode header value for Kiro IDE requests	kiroIDEAgentMode = "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) // endpointAliases maps user preference values to canonical endpoint names.var endpointAliases = map[string]string{	"codewhisperer": "codewhisperer",	"ide":           "codewhisperer",	"amazonq":       "amazonq",	"q":             "amazonq",	"cli":           "amazonq",} func enqueueTranslatedSSE(out chan<- cliproxyexecutor.StreamChunk, chunk []byte) {	if len(chunk) == 0 {		return	}	out <- cliproxyexecutor.StreamChunk{Payload: append(bytes.Clone(chunk), '\n', '\n')}} // 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",		},	}} // resolveKiroAPIRegion determines the AWS region for Kiro API calls.// 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 resolveKiroAPIRegion(auth *cliproxyauth.Auth) string {	if auth == nil || auth.Metadata == nil {		return kiroDefaultRegion	}	// Priority 1: Explicit api_region override	if r, ok := auth.Metadata["api_region"].(string); ok && r != "" {		log.Debugf("kiro: using region %s (source: api_region)", r)		return r	}	// Priority 2: Extract from ProfileARN	if profileArn, ok := auth.Metadata["profile_arn"].(string); ok && profileArn != "" {		if arnRegion := extractRegionFromProfileARN(profileArn); arnRegion != "" {			log.Debugf("kiro: using region %s (source: profile_arn)", arnRegion)			return arnRegion		}	}	// 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: default)", kiroDefaultRegion)	return kiroDefaultRegion} // 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	} 	region := resolveKiroAPIRegion(auth)	log.Debugf("kiro: using region %s", region) 	configs := buildKiroEndpointConfigs(region) 	preference := getAuthValue(auth, "preferred_endpoint")	if preference == "" {		return configs	} 	targetName, ok := endpointAliases[preference]	if !ok {		return configs	} 	var preferred, others []kiroEndpointConfig	for _, cfg := range configs {		if strings.ToLower(cfg.Name) == targetName {			preferred = append(preferred, cfg)		} else {			others = append(others, cfg)		}	} 	if len(preferred) == 0 {		return configs	}	return append(preferred, others...)} // 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	profileArnMu sync.Mutex // Serializes profileArn fetches to prevent concurrent map writes} // buildKiroPayloadForFormat builds the Kiro API payload based on the source format.// This is critical because OpenAI and Claude formats have different tool structures:// - OpenAI: tools[].function.name, tools[].function.description// - Claude: tools[].name, tools[].description// headers parameter allows checking Anthropic-Beta header for thinking mode detection.// Returns the serialized JSON payload and a boolean indicating whether thinking mode was injected.func buildKiroPayloadForFormat(body []byte, modelID, profileArn, origin string, isAgentic, isChatOnly bool, sourceFormat sdktranslator.Format, headers http.Header) ([]byte, bool) {	switch sourceFormat.String() {	case "openai":		log.Debugf("kiro: using OpenAI payload builder for source format: %s", sourceFormat.String())		return kiroopenai.BuildKiroPayloadFromOpenAI(body, modelID, profileArn, origin, isAgentic, isChatOnly, headers, nil)	case "kiro":		// Body is already in Kiro format — pass through directly		log.Debugf("kiro: body already in Kiro format, passing through directly")		return body, false	default:		// Default to Claude format		log.Debugf("kiro: using Claude payload builder for source format: %s", sourceFormat.String())		return kiroclaude.BuildKiroPayload(body, modelID, profileArn, origin, isAgentic, isChatOnly, headers, nil)	}} // NewKiroExecutor creates a new Kiro executor instance.func NewKiroExecutor(cfg *config.Config) *KiroExecutor {	return &KiroExecutor{cfg: cfg}} // Identifier returns the unique identifier for this executor.func (e *KiroExecutor) Identifier() string { return "kiro" } // applyDynamicFingerprint applies account-specific fingerprint headers to the request.func applyDynamicFingerprint(req *http.Request, auth *cliproxyauth.Auth) {	accountKey := getAccountKey(auth)	fp := kiroauth.GlobalFingerprintManager().GetFingerprint(accountKey) 

Lines 1–500 of 4706. Load more