package executor import (	"bufio"	"bytes"	"context"	"encoding/base64"	"encoding/binary"	"encoding/json"	"fmt"	"io"	"net/http"	"regexp"	"strings"	"sync"	"time"	"unicode/utf8" 	"github.com/google/uuid"	kiroauth "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/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"	"github.com/tidwall/gjson" 	"github.com/gin-gonic/gin") const (	// Kiro API common constants	kiroContentType    = "application/x-amz-json-1.0"	kiroAcceptStream   = "*/*"	kiroMaxMessageSize = 10 * 1024 * 1024 // 10MB max message size for event stream	kiroMaxToolDescLen = 10237            // Kiro API limit is 10240 bytes, leave room for "..."	// kiroUserAgent matches amq2api format 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 matching amq2api	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" 	// Thinking mode support - based on amq2api implementation	// These tags wrap reasoning content in the response stream	thinkingStartTag = "<thinking>"	thinkingEndTag   = "</thinking>"	// thinkingHint is injected into the request to enable interleaved thinking mode	// This tells the model to use thinking tags and sets the max thinking length	thinkingHint     = "<thinking_mode>interleaved</thinking_mode><max_thinking_length>16000</max_thinking_length>" 	// kiroAgenticSystemPrompt is injected only for -agentic models to prevent timeouts on large writes.	// AWS Kiro API has a 2-3 minute timeout for large file write operations.	kiroAgenticSystemPrompt = `# CRITICAL: CHUNKED WRITE PROTOCOL (MANDATORY) You MUST follow these rules for ALL file operations. Violation causes server timeouts and task failure. ## ABSOLUTE LIMITS- **MAXIMUM 350 LINES** per single write/edit operation - NO EXCEPTIONS- **RECOMMENDED 300 LINES** or less for optimal performance- **NEVER** write entire files in one operation if >300 lines ## MANDATORY CHUNKED WRITE STRATEGY ### For NEW FILES (>300 lines total):1. FIRST: Write initial chunk (first 250-300 lines) using write_to_file/fsWrite2. THEN: Append remaining content in 250-300 line chunks using file append operations3. REPEAT: Continue appending until complete ### For EDITING EXISTING FILES:1. Use surgical edits (apply_diff/targeted edits) - change ONLY what's needed2. NEVER rewrite entire files - use incremental modifications3. Split large refactors into multiple small, focused edits ### For LARGE CODE GENERATION:1. Generate in logical sections (imports, types, functions separately)2. Write each section as a separate operation3. Use append operations for subsequent sections ## EXAMPLES OF CORRECT BEHAVIOR ✅ CORRECT: Writing a 600-line file- Operation 1: Write lines 1-300 (initial file creation)- Operation 2: Append lines 301-600 ✅ CORRECT: Editing multiple functions- Operation 1: Edit function A- Operation 2: Edit function B- Operation 3: Edit function C ❌ WRONG: Writing 500 lines in single operation → TIMEOUT❌ WRONG: Rewriting entire file to change 5 lines → TIMEOUT❌ WRONG: Generating massive code blocks without chunking → TIMEOUT ## WHY THIS MATTERS- Server has 2-3 minute timeout for operations- Large writes exceed timeout and FAIL completely- Chunked writes are FASTER and more RELIABLE- Failed writes waste time and require retry REMEMBER: When in doubt, write LESS per operation. Multiple small operations > one large operation.`) // 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.func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig {	if auth == nil {		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} // 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" } // PrepareRequest prepares the HTTP request before execution.func (e *KiroExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil }  // 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) {	accessToken, profileArn := kiroCredentials(auth)	if accessToken == "" {		return resp, fmt.Errorf("kiro: access token not found in auth")	}	if profileArn == "" {		// Only warn if not using builder-id auth (which doesn't need profileArn)		if auth == nil || auth.Metadata == nil {			log.Debugf("kiro: profile ARN not found in auth (may be normal for builder-id)")		} else if authMethod, ok := auth.Metadata["auth_method"].(string); !ok || authMethod != "builder-id" {			log.Warnf("kiro: profile ARN not found in auth, API calls may fail")		}	} 	reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)	defer reporter.trackFailure(ctx, &err) 	// 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			accessToken, profileArn = kiroCredentials(auth)			log.Infof("kiro: token refreshed successfully before request")		}	} 	from := opts.SourceFormat	to := sdktranslator.FromString("kiro")	body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true) 	kiroModelID := e.mapModelToKiro(req.Model)		// Check if this is an agentic model variant	isAgentic := strings.HasSuffix(req.Model, "-agentic")		// Check if this is a chat-only model variant (no tool calling)	isChatOnly := strings.HasSuffix(req.Model, "-chat")		// Determine initial origin - always use AI_EDITOR to match AIClient-2-API behavior	// AIClient-2-API uses AI_EDITOR for all models, which is the Kiro IDE quota	// Note: CLI origin is for Amazon Q quota, but AIClient-2-API doesn't use it	currentOrigin := "AI_EDITOR"		// Determine if profileArn should be included based on auth method	// profileArn is only needed for social auth (Google OAuth), not for builder-id (AWS SSO)	effectiveProfileArn := profileArn	if auth != nil && auth.Metadata != nil {		if authMethod, ok := auth.Metadata["auth_method"].(string); ok && authMethod == "builder-id" {			effectiveProfileArn = "" // Don't include profileArn for builder-id auth		}	}		kiroPayload := e.buildKiroPayload(body, kiroModelID, effectiveProfileArn, currentOrigin, isAgentic, isChatOnly) 	// Execute with retry on 401/403 and 429 (quota exhausted)	resp, err = e.executeWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, kiroPayload, body, from, to, reporter, currentOrigin, kiroModelID, isAgentic, isChatOnly)	return resp, err} // executeWithRetry performs the actual HTTP request with automatic retry on auth errors.// Supports automatic fallback between endpoints with different quotas:// - 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) {	var resp cliproxyexecutor.Response	maxRetries := 2 // Allow retries for token refresh + endpoint fallback	endpointConfigs := getKiroEndpointConfigs(auth) 	for endpointIdx := 0; endpointIdx < len(endpointConfigs); endpointIdx++ {		endpointConfig := endpointConfigs[endpointIdx]		url := endpointConfig.URL		// Use this endpoint's compatible Origin (critical for avoiding 403 errors)		currentOrigin = endpointConfig.Origin				// Rebuild payload with the correct origin for this endpoint		// Each endpoint requires its matching Origin value in the request body		kiroPayload = e.buildKiroPayload(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly)				log.Debugf("kiro: trying endpoint %d/%d: %s (Name: %s, Origin: %s)",			endpointIdx+1, len(endpointConfigs), url, endpointConfig.Name, currentOrigin) 	for attempt := 0; attempt <= maxRetries; attempt++ {		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroPayload))		if err != nil {			return resp, err		} 		httpReq.Header.Set("Content-Type", kiroContentType)		httpReq.Header.Set("Authorization", "Bearer "+accessToken)		httpReq.Header.Set("Accept", kiroAcceptStream)		// Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors)		httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget)		httpReq.Header.Set("User-Agent", kiroUserAgent)		httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)		httpReq.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")		httpReq.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String()) 		var attrs map[string]string		if auth != nil {			attrs = auth.Attributes		}		util.ApplyCustomHeadersFromAttrs(httpReq, attrs) 		var authID, authLabel, authType, authValue string		if auth != nil {			authID = auth.ID			authLabel = auth.Label			authType, authValue = auth.AccountInfo()		}		recordAPIRequest(ctx, e.cfg, upstreamRequestLog{			URL:       url,			Method:    http.MethodPost,			Headers:   httpReq.Header.Clone(),			Body:      kiroPayload,			Provider:  e.Identifier(),			AuthID:    authID,			AuthLabel: authLabel,			AuthType:  authType,			AuthValue: authValue,		}) 		httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 120*time.Second)		httpResp, err := httpClient.Do(httpReq)		if err != nil {			recordAPIResponseError(ctx, e.cfg, err)			return resp, err		}		recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) 		// Handle 429 errors (quota exhausted) - try next endpoint		// Each endpoint has its own quota pool, so we can try different endpoints		if httpResp.StatusCode == 429 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			log.Warnf("kiro: %s endpoint quota exhausted (429), will try next endpoint", endpointConfig.Name)						// Break inner retry loop to try next endpoint (which has different quota)			break		} 		// Handle 5xx server errors with exponential backoff retry		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)				backoff := time.Duration(1<<attempt) * time.Second				if backoff > 30*time.Second {					backoff = 30 * time.Second				}				log.Warnf("kiro: server error %d, retrying in %v (attempt %d/%d)", httpResp.StatusCode, backoff, attempt+1, maxRetries)				time.Sleep(backoff)				continue			}			log.Errorf("kiro: server error %d after %d retries", httpResp.StatusCode, maxRetries)			return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}		} 		// Handle 401 errors with token refresh and retry		// 401 = Unauthorized (token expired/invalid) - refresh token		if httpResp.StatusCode == 401 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			if attempt < maxRetries {				log.Warnf("kiro: received 401 error, attempting token refresh and retry (attempt %d/%d)", attempt+1, maxRetries+1) 				refreshedAuth, refreshErr := e.Refresh(ctx, auth)				if refreshErr != nil {					log.Errorf("kiro: token refresh failed: %v", refreshErr)					return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}				} 				if refreshedAuth != nil {					auth = refreshedAuth					accessToken, profileArn = kiroCredentials(auth)					// Rebuild payload with new profile ARN if changed					kiroPayload = e.buildKiroPayload(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly)					log.Infof("kiro: token refreshed successfully, retrying request")					continue				}			} 			log.Warnf("kiro request error, status: 401, body: %s", summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))			return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}		} 		// Handle 402 errors - Monthly Limit Reached		if httpResp.StatusCode == 402 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			log.Warnf("kiro: received 402 (monthly limit). Upstream body: %s", string(respBody)) 			// Return upstream error body directly			return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}		} 		// Handle 403 errors - Access Denied / Token Expired		// Do NOT switch endpoints for 403 errors		if httpResp.StatusCode == 403 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			// Log the 403 error details for debugging			log.Warnf("kiro: received 403 error (attempt %d/%d), body: %s", attempt+1, maxRetries+1, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody)) 			respBodyStr := string(respBody) 			// 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")				return resp, statusErr{code: httpResp.StatusCode, msg: "account suspended: " + string(respBody)}			} 			// Check if this looks like a token-related 403 (some APIs return 403 for expired tokens)			isTokenRelated := strings.Contains(respBodyStr, "token") ||				strings.Contains(respBodyStr, "expired") ||				strings.Contains(respBodyStr, "invalid") ||				strings.Contains(respBodyStr, "unauthorized") 			if isTokenRelated && attempt < maxRetries {				log.Warnf("kiro: 403 appears token-related, attempting token refresh")				refreshedAuth, refreshErr := e.Refresh(ctx, auth)				if refreshErr != nil {					log.Errorf("kiro: token refresh failed: %v", refreshErr)					// Token refresh failed - return error immediately					return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}				}				if refreshedAuth != nil {					auth = refreshedAuth					accessToken, profileArn = kiroCredentials(auth)					kiroPayload = e.buildKiroPayload(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly)					log.Infof("kiro: token refreshed for 403, retrying request")					continue				}			} 			// For non-token 403 or after max retries, return error immediately			// Do NOT switch endpoints for 403 errors			log.Warnf("kiro: 403 error, returning immediately (no endpoint switch)")			return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}		} 		if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {			b, _ := io.ReadAll(httpResp.Body)			appendAPIResponseChunk(ctx, e.cfg, b)			log.Debugf("kiro request error, status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))			err = statusErr{code: httpResp.StatusCode, msg: string(b)}			if errClose := httpResp.Body.Close(); errClose != nil {				log.Errorf("response body close error: %v", errClose)			}			return resp, err		} 		defer func() {			if errClose := httpResp.Body.Close(); errClose != nil {				log.Errorf("response body close error: %v", errClose)			}		}() 		content, toolUses, usageInfo, err := e.parseEventStream(httpResp.Body)		if err != nil {			recordAPIResponseError(ctx, e.cfg, err)

Lines 1–500 of 3696. Load more