diff --git a/internal/runtime/executor/kiro_executor.go b/internal/runtime/executor/kiro_executor.go index 7c10d03e..2c862807 100644 --- a/internal/runtime/executor/kiro_executor.go +++ b/internal/runtime/executor/kiro_executor.go @@ -334,12 +334,30 @@ type kiroEndpointConfig struct { 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. +// kiroDefaultRegion is the default AWS region for Kiro API endpoints. +// Used when no region is specified in auth metadata. +const kiroDefaultRegion = "us-east-1" + +// extractRegionFromProfileARN extracts the AWS region from a ProfileARN. +// ARN format: arn:aws:codewhisperer:REGION:ACCOUNT:profile/PROFILE_ID +// Returns empty string if region cannot be extracted. +func extractRegionFromProfileARN(profileArn string) string { + if profileArn == "" { + return "" + } + parts := strings.Split(profileArn, ":") + if len(parts) >= 4 && parts[3] != "" { + return parts[3] + } + return "" +} + +// buildKiroEndpointConfigs creates endpoint configurations for the specified region. +// This enables dynamic region support for Enterprise/IdC users in non-us-east-1 regions. // // 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 +// - CodeWhisperer endpoint (codewhisperer.{region}.amazonaws.com): Uses AI_EDITOR origin and AmazonCodeWhispererStreamingService target +// - Amazon Q endpoint (q.{region}.amazonaws.com): Uses CLI origin and AmazonQDeveloperStreamingService target // // Mismatched combinations will result in 403 Forbidden errors. // @@ -348,38 +366,86 @@ type kiroEndpointConfig struct { // 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", - }, +func buildKiroEndpointConfigs(region string) []kiroEndpointConfig { + if region == "" { + region = kiroDefaultRegion + } + return []kiroEndpointConfig{ + { + URL: fmt.Sprintf("https://codewhisperer.%s.amazonaws.com/generateAssistantResponse", region), + Origin: "AI_EDITOR", + AmzTarget: "AmazonCodeWhispererStreamingService.GenerateAssistantResponse", + Name: "CodeWhisperer", + }, + { + URL: fmt.Sprintf("https://q.%s.amazonaws.com/", region), + Origin: "CLI", + AmzTarget: "AmazonQDeveloperStreamingService.SendMessage", + Name: "AmazonQ", + }, + } } +// kiroEndpointConfigs is kept for backward compatibility with default us-east-1 region. +// Prefer using buildKiroEndpointConfigs(region) for dynamic region support. +var kiroEndpointConfigs = buildKiroEndpointConfigs(kiroDefaultRegion) + // getKiroEndpointConfigs returns the list of Kiro API endpoint configurations to try in order. +// Supports dynamic region based on auth metadata "api_region", "profile_arn", or "region" field. // Supports reordering based on "preferred_endpoint" in auth metadata/attributes. // For IDC auth method, automatically uses CodeWhisperer endpoint with CLI origin. +// +// Region priority: +// 1. auth.Metadata["api_region"] - explicit API region override +// 2. ProfileARN region - extracted from arn:aws:service:REGION:account:resource +// 3. auth.Metadata["region"] - OIDC/Identity region (may differ from API region) +// 4. kiroDefaultRegion (us-east-1) - fallback func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig { if auth == nil { return kiroEndpointConfigs } + // Determine API region with priority: api_region > profile_arn > region > default + region := kiroDefaultRegion + regionSource := "default" + + if auth.Metadata != nil { + // Priority 1: Explicit api_region override + if r, ok := auth.Metadata["api_region"].(string); ok && r != "" { + region = r + regionSource = "api_region" + } else { + // Priority 2: Extract from ProfileARN + if profileArn, ok := auth.Metadata["profile_arn"].(string); ok && profileArn != "" { + if arnRegion := extractRegionFromProfileARN(profileArn); arnRegion != "" { + region = arnRegion + regionSource = "profile_arn" + } + } + // Priority 3: OIDC region (only if not already set from profile_arn) + if regionSource == "default" { + if r, ok := auth.Metadata["region"].(string); ok && r != "" { + region = r + regionSource = "region" + } + } + } + } + + log.Debugf("kiro: using region %s (source: %s)", region, regionSource) + + // Build endpoint configs for the specified region + endpointConfigs := buildKiroEndpointConfigs(region) + // 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 + if strings.ToLower(authMethod) == "idc" { + log.Debugf("kiro: IDC auth, using CodeWhisperer endpoint (region: %s)", region) + return endpointConfigs } } @@ -396,7 +462,7 @@ func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig { } if preference == "" { - return kiroEndpointConfigs + return endpointConfigs } preference = strings.ToLower(strings.TrimSpace(preference)) @@ -405,7 +471,7 @@ func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig { var sorted []kiroEndpointConfig var remaining []kiroEndpointConfig - for _, cfg := range kiroEndpointConfigs { + for _, cfg := range endpointConfigs { name := strings.ToLower(cfg.Name) // Check for matches // CodeWhisperer aliases: codewhisperer, ide @@ -426,7 +492,7 @@ func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig { // If preference didn't match anything, return default if len(sorted) == 0 { - return kiroEndpointConfigs + return endpointConfigs } // Combine: preferred first, then others @@ -445,7 +511,7 @@ func isIDCAuth(auth *cliproxyauth.Auth) bool { return false } authMethod, _ := auth.Metadata["auth_method"].(string) - return authMethod == "idc" + return strings.ToLower(authMethod) == "idc" } // buildKiroPayloadForFormat builds the Kiro API payload based on the source format.