package executor import (	"bufio"	"bytes"	"context"	"encoding/base64"	"encoding/binary"	"encoding/json"	"fmt"	"io"	"net/http"	"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	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") // 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.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} // 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(_ *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")	} 	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) 	// 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("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) 				// 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						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 {					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, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)						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, stopReason, 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 := getTokenizer(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 := getTokenizer(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			// stopReason is extracted from upstream response by parseEventStream			kiroResponse := kiroclaude.BuildClaudeResponse(content, toolUses, req.Model, usageInfo, stopReason)			out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, kiroResponse, nil)			resp = cliproxyexecutor.Response{Payload: []byte(out)}			return resp, nil		}		// Inner retry loop exhausted for this endpoint, try next endpoint		// Note: This code is unreachable because all paths in the inner loop		// either return or continue. Kept as comment for documentation.	} 	// All endpoints exhausted	if last429Err != nil {		return resp, last429Err	}	return resp, fmt.Errorf("kiro: all endpoints exhausted")}

Lines 1–500 of 3157. Load more