package executor import (	"bufio"	"bytes"	"context"	"encoding/base64"	"encoding/binary"	"encoding/json"	"fmt"	"io"	"net/http"	"os"	"path/filepath"	"strings"	"sync"	"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") // 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) // 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" } // PrepareRequest prepares the HTTP request before execution.func (e *KiroExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error {	if req == nil {		return nil	}	accessToken, _ := kiroCredentials(auth)	if strings.TrimSpace(accessToken) == "" {		return statusErr{code: http.StatusUnauthorized, msg: "missing access token"}	}	if isIDCAuth(auth) {		req.Header.Set("User-Agent", kiroIDEUserAgent)		req.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent)		req.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec)	} else {		req.Header.Set("User-Agent", kiroUserAgent)		req.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)	}	req.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")	req.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String())	req.Header.Set("Authorization", "Bearer "+accessToken)	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(req, attrs)	return nil} // HttpRequest injects Kiro credentials into the request and executes it.func (e *KiroExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) {	if req == nil {		return nil, fmt.Errorf("kiro executor: request is nil")	}	if ctx == nil {		ctx = req.Context()	}	httpReq := req.WithContext(ctx)	if errPrepare := e.PrepareRequest(httpReq, auth); errPrepare != nil {		return nil, errPrepare	}	httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)	return httpClient.Do(httpReq)} // 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")	} 	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			// Persist the refreshed auth to file so subsequent requests use it			if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {				log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)			}			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) 	// Determine agentic mode and effective profile ARN using helper functions	isAgentic, isChatOnly := determineAgenticMode(req.Model)	effectiveProfileArn := getEffectiveProfileArnWithWarning(auth, profileArn) 	// Execute with retry on 401/403 and 429 (quota exhausted)	// Note: currentOrigin and kiroPayload are built inside executeWithRetry for each endpoint	resp, err = e.executeWithRetry(ctx, auth, req, opts, accessToken, effectiveProfileArn, nil, body, from, to, reporter, "", 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)	var last429Err error 	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, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers) 		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("Accept", kiroAcceptStream)			// Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors)			httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget) 			// Use different headers based on auth type			// IDC auth uses Kiro IDE style headers (from kiro2api)			// Other auth types use Amazon Q CLI style headers			if isIDCAuth(auth) {				httpReq.Header.Set("User-Agent", kiroIDEUserAgent)				httpReq.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent)				httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec)				log.Debugf("kiro: using Kiro IDE headers for IDC auth")			} else {				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()) 			// Bearer token authentication for all auth types (Builder ID, IDC, social, etc.)			httpReq.Header.Set("Authorization", "Bearer "+accessToken) 			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) 				// Preserve last 429 so callers can correctly backoff when all endpoints are exhausted				last429Err = statusErr{code: httpResp.StatusCode, msg: string(respBody)} 				log.Warnf("kiro: %s endpoint quota exhausted (429), will try next endpoint, body: %s",					endpointConfig.Name, summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody)) 				// 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						// Persist the refreshed auth to file so subsequent requests use it						if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {							log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)							// Continue anyway - the token is valid for this request						}						accessToken, profileArn = kiroCredentials(auth)						// Rebuild payload with new profile ARN if changed						kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)						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 {

Lines 1–500 of 3341. Load more