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 (	// kiroEndpoint is the Amazon Q streaming endpoint for chat API (GenerateAssistantResponse).	// Note: This is different from the CodeWhisperer management endpoint (codewhisperer.us-east-1.amazonaws.com)	// used in aws_auth.go for GetUsageLimits, ListProfiles, etc. Both endpoints are correct	// for their respective API operations.	kiroEndpoint       = "https://q.us-east-1.amazonaws.com"	kiroTargetChat     = "AmazonCodeWhispererStreamingService.GenerateAssistantResponse"	kiroContentType    = "application/x-amz-json-1.0"	kiroAcceptStream   = "application/vnd.amazon.eventstream"	kiroMaxMessageSize = 10 * 1024 * 1024 // 10MB max message size for event stream	kiroMaxToolDescLen = 10237            // Kiro API limit is 10240 bytes, leave room for "..." 	// 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.`) // 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 based on model type	// Opus models use AI_EDITOR (Kiro IDE quota), others start with CLI (Amazon Q quota)	var currentOrigin string	if strings.Contains(strings.ToLower(req.Model), "opus") {		currentOrigin = "AI_EDITOR"	} else {		currentOrigin = "CLI"	}		// 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 from CLI (Amazon Q) quota to AI_EDITOR (Kiro IDE) quota on 429.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 + origin fallback 	for attempt := 0; attempt <= maxRetries; attempt++ {		url := kiroEndpoint		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("x-amz-target", kiroTargetChat)		httpReq.Header.Set("Authorization", "Bearer "+accessToken)		httpReq.Header.Set("Accept", kiroAcceptStream) 		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) with origin fallback		if httpResp.StatusCode == 429 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			// If currently using CLI quota and it's exhausted, switch to AI_EDITOR (Kiro IDE) quota			if currentOrigin == "CLI" {				log.Warnf("kiro: Amazon Q (CLI) quota exhausted (429), switching to Kiro (AI_EDITOR) fallback")				currentOrigin = "AI_EDITOR" 				// Rebuild payload with new origin				kiroPayload = e.buildKiroPayload(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly) 				// Retry with new origin				continue			} 			// Already on AI_EDITOR or other origin, return the error			log.Debugf("kiro request error, status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))			return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)}		} 		// Handle 401/403 errors with token refresh and retry		if httpResp.StatusCode == 401 || httpResp.StatusCode == 403 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			if attempt < maxRetries {				log.Warnf("kiro: received %d error, attempting token refresh and retry (attempt %d/%d)", httpResp.StatusCode, 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.Debugf("kiro request error, status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody))			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)			return resp, err		} 		// Fallback for usage if missing from upstream		if usageInfo.TotalTokens == 0 {			if enc, encErr := tokenizerForModel(req.Model); encErr == nil {				if inp, countErr := countOpenAIChatTokens(enc, opts.OriginalRequest); countErr == nil {					usageInfo.InputTokens = inp				}			}			if len(content) > 0 {				// Use tiktoken for more accurate output token calculation				if enc, encErr := tokenizerForModel(req.Model); encErr == nil {					if tokenCount, countErr := enc.Count(content); countErr == nil {						usageInfo.OutputTokens = int64(tokenCount)					}				}				// Fallback to character count estimation if tiktoken fails				if usageInfo.OutputTokens == 0 {					usageInfo.OutputTokens = int64(len(content) / 4)					if usageInfo.OutputTokens == 0 {						usageInfo.OutputTokens = 1					}				}			}			usageInfo.TotalTokens = usageInfo.InputTokens + usageInfo.OutputTokens		} 		appendAPIResponseChunk(ctx, e.cfg, []byte(content))		reporter.publish(ctx, usageInfo) 		// Build response in Claude format for Kiro translator		kiroResponse := e.buildClaudeResponse(content, toolUses, req.Model, usageInfo)		out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, kiroResponse, nil)		resp = cliproxyexecutor.Response{Payload: []byte(out)}		return resp, nil	} 	return resp, fmt.Errorf("kiro: max retries exceeded")} // ExecuteStream handles streaming requests to Kiro API.// Supports automatic token refresh on 401/403 errors and quota fallback on 429.func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {	accessToken, profileArn := kiroCredentials(auth)	if accessToken == "" {		return nil, 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 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			accessToken, profileArn = kiroCredentials(auth)			log.Infof("kiro: token refreshed successfully before stream 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 based on model type	// Opus models use AI_EDITOR (Kiro IDE quota), others start with CLI (Amazon Q quota)	var currentOrigin string	if strings.Contains(strings.ToLower(req.Model), "opus") {		currentOrigin = "AI_EDITOR"	} else {		currentOrigin = "CLI"	}		// 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 stream with retry on 401/403 and 429 (quota exhausted)	return e.executeStreamWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, kiroPayload, body, from, reporter, currentOrigin, kiroModelID, isAgentic, isChatOnly)} // executeStreamWithRetry performs the streaming HTTP request with automatic retry on auth errors.// Supports automatic fallback from CLI (Amazon Q) quota to AI_EDITOR (Kiro IDE) quota on 429.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) {	maxRetries := 2 // Allow retries for token refresh + origin fallback 	for attempt := 0; attempt <= maxRetries; attempt++ {		url := kiroEndpoint		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroPayload))		if err != nil {			return nil, err		} 		httpReq.Header.Set("Content-Type", kiroContentType)		httpReq.Header.Set("x-amz-target", kiroTargetChat)		httpReq.Header.Set("Authorization", "Bearer "+accessToken)		httpReq.Header.Set("Accept", kiroAcceptStream) 		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, 0)		httpResp, err := httpClient.Do(httpReq)		if err != nil {			recordAPIResponseError(ctx, e.cfg, err)			return nil, err		}		recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) 		// Handle 429 errors (quota exhausted) with origin fallback		if httpResp.StatusCode == 429 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			// If currently using CLI quota and it's exhausted, switch to AI_EDITOR (Kiro IDE) quota			if currentOrigin == "CLI" {				log.Warnf("kiro: stream Amazon Q (CLI) quota exhausted (429), switching to Kiro (AI_EDITOR) fallback")				currentOrigin = "AI_EDITOR" 				// Rebuild payload with new origin				kiroPayload = e.buildKiroPayload(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly) 				// Retry with new origin				continue			} 			// Already on AI_EDITOR or other origin, return the error			log.Debugf("kiro stream error, status: %d, body: %s", httpResp.StatusCode, string(respBody))			return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}		} 		// Handle 401/403 errors with token refresh and retry		if httpResp.StatusCode == 401 || httpResp.StatusCode == 403 {			respBody, _ := io.ReadAll(httpResp.Body)			_ = httpResp.Body.Close()			appendAPIResponseChunk(ctx, e.cfg, respBody) 			if attempt < maxRetries {				log.Warnf("kiro: stream received %d error, attempting token refresh and retry (attempt %d/%d)", httpResp.StatusCode, attempt+1, maxRetries+1) 				refreshedAuth, refreshErr := e.Refresh(ctx, auth)				if refreshErr != nil {					log.Errorf("kiro: token refresh failed: %v", refreshErr)					return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)}

Lines 1–500 of 3006. Load more