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/x-amz-json-1.0"	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 amq2api format for User-Agent header (Amazon Q CLI style)	kiroUserAgent = "aws-sdk-rust/1.3.9 os/macos lang/rust/1.87.0"	// kiroFullUserAgent is the complete x-amz-user-agent header matching amq2api (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 (from kiro2api - for IDC auth)	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" 	// 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 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} // kiroEndpointConfigs defines the available Kiro API endpoints with their compatible configurations.// The order determines fallback priority: primary endpoint first, then fallbacks.//// CRITICAL: Each endpoint MUST use its compatible Origin and AmzTarget values:// - CodeWhisperer endpoint (codewhisperer.us-east-1.amazonaws.com): Uses AI_EDITOR origin and AmazonCodeWhispererStreamingService target// - Amazon Q endpoint (q.us-east-1.amazonaws.com): Uses CLI origin and AmazonQDeveloperStreamingService target//// Mismatched combinations will result in 403 Forbidden errors.//// NOTE: CodeWhisperer is set as the default endpoint because:// 1. Most tokens come from Kiro IDE / VSCode extensions (AWS Builder ID auth)// 2. These tokens use AI_EDITOR origin which is only compatible with CodeWhisperer endpoint// 3. Amazon Q endpoint requires CLI origin which is for Amazon Q CLI tokens// This matches the AIClient-2-API-main project's configuration.var kiroEndpointConfigs = []kiroEndpointConfig{	{		URL:       "https://codewhisperer.us-east-1.amazonaws.com/generateAssistantResponse",		Origin:    "AI_EDITOR",		AmzTarget: "AmazonCodeWhispererStreamingService.GenerateAssistantResponse",		Name:      "CodeWhisperer",	},	{		URL:       "https://q.us-east-1.amazonaws.com/",		Origin:    "CLI",		AmzTarget: "AmazonQDeveloperStreamingService.SendMessage",		Name:      "AmazonQ",	},} // getKiroEndpointConfigs returns the list of Kiro API endpoint configurations to try in order.// Supports reordering based on "preferred_endpoint" in auth metadata/attributes.// For IDC auth method, automatically uses CodeWhisperer endpoint with CLI origin.func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig {	if auth == nil {		return kiroEndpointConfigs	} 	// For IDC auth, use CodeWhisperer endpoint with AI_EDITOR origin (same as Social auth)	// Based on kiro2api analysis: IDC tokens work with CodeWhisperer 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 authMethod == "idc" {			log.Debugf("kiro: IDC auth, using CodeWhisperer endpoint")			return kiroEndpointConfigs		}	} 	// 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 kiroEndpointConfigs	} 	preference = strings.ToLower(strings.TrimSpace(preference)) 	// Create new slice to avoid modifying global state	var sorted []kiroEndpointConfig	var remaining []kiroEndpointConfig 	for _, cfg := range kiroEndpointConfigs {		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 kiroEndpointConfigs	} 	// 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} // isIDCAuth checks if the auth uses IDC (Identity Center) authentication method.func isIDCAuth(auth *cliproxyauth.Auth) bool {	if auth == nil || auth.Metadata == nil {		return false	}	authMethod, _ := auth.Metadata["auth_method"].(string)	return authMethod == "idc"} // 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)	default:		// Default to Claude format (also handles "claude", "kiro", etc.)		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 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 headersfunc 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.

Lines 1–500 of 3891. Load more