CLIProxyAPIPlus/internal/auth/kiro/codewhisperer_client.go

// Package kiro provides CodeWhisperer API client for fetching user info.
package kiro

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
	log "github.com/sirupsen/logrus"
)

// CodeWhispererClient handles CodeWhisperer API calls.
type CodeWhispererClient struct {
	httpClient *http.Client
}

// UsageLimitsResponse represents the getUsageLimits API response.
type UsageLimitsResponse struct {
	DaysUntilReset     *int              `json:"daysUntilReset,omitempty"`
	NextDateReset      *float64          `json:"nextDateReset,omitempty"`
	UserInfo           *UserInfo         `json:"userInfo,omitempty"`
	SubscriptionInfo   *SubscriptionInfo `json:"subscriptionInfo,omitempty"`
	UsageBreakdownList []UsageBreakdown  `json:"usageBreakdownList,omitempty"`
}

// UserInfo contains user information from the API.
type UserInfo struct {
	Email  string `json:"email,omitempty"`
	UserID string `json:"userId,omitempty"`
}

// SubscriptionInfo contains subscription details.
type SubscriptionInfo struct {
	SubscriptionTitle string `json:"subscriptionTitle,omitempty"`
	Type              string `json:"type,omitempty"`
}

// UsageBreakdown contains usage details.
type UsageBreakdown struct {
	UsageLimit                *int     `json:"usageLimit,omitempty"`
	CurrentUsage              *int     `json:"currentUsage,omitempty"`
	UsageLimitWithPrecision   *float64 `json:"usageLimitWithPrecision,omitempty"`
	CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision,omitempty"`
	NextDateReset             *float64 `json:"nextDateReset,omitempty"`
	DisplayName               string   `json:"displayName,omitempty"`
	ResourceType              string   `json:"resourceType,omitempty"`
}

// NewCodeWhispererClient creates a new CodeWhisperer client.
func NewCodeWhispererClient(cfg *config.Config, machineID string) *CodeWhispererClient {
	client := &http.Client{Timeout: 30 * time.Second}
	if cfg != nil {
		client = util.SetProxy(&cfg.SDKConfig, client)
	}
	return &CodeWhispererClient{
		httpClient: client,
	}
}

// GetUsageLimits fetches usage limits and user info from CodeWhisperer API.
// This is the recommended way to get user email after login.
func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken, clientID, refreshToken, profileArn string) (*UsageLimitsResponse, error) {
	queryParams := map[string]string{
		"origin":       "AI_EDITOR",
		"resourceType": "AGENTIC_REQUEST",
	}
	// Determine endpoint based on profileArn region
	endpoint := GetKiroAPIEndpointFromProfileArn(profileArn)
	if profileArn != "" {
		queryParams["profileArn"] = profileArn
	} else {
		queryParams["isEmailRequired"] = "true"
	}
	url := buildURL(endpoint, pathGetUsageLimits, queryParams)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	accountKey := GetAccountKey(clientID, refreshToken)
	setRuntimeHeaders(req, accessToken, accountKey)

	log.Debugf("codewhisperer: GET %s", url)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	log.Debugf("codewhisperer: status=%d, body=%s", resp.StatusCode, string(body))

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(body))
	}

	var result UsageLimitsResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return &result, nil
}

// FetchUserEmailFromAPI fetches user email using CodeWhisperer getUsageLimits API.
// This is more reliable than JWT parsing as it uses the official API.
func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessToken, clientID, refreshToken string) string {
	resp, err := c.GetUsageLimits(ctx, accessToken, clientID, refreshToken, "")
	if err != nil {
		log.Debugf("codewhisperer: failed to get usage limits: %v", err)
		return ""
	}

	if resp.UserInfo != nil && resp.UserInfo.Email != "" {
		log.Debugf("codewhisperer: got email from API: %s", resp.UserInfo.Email)
		return resp.UserInfo.Email
	}

	log.Debugf("codewhisperer: no email in response")
	return ""
}

// FetchUserEmailWithFallback fetches user email with multiple fallback methods.
// Priority: 1. CodeWhisperer API  2. userinfo endpoint  3. JWT parsing
func FetchUserEmailWithFallback(ctx context.Context, cfg *config.Config, accessToken, clientID, refreshToken string) string {
	// Method 1: Try CodeWhisperer API (most reliable)
	cwClient := NewCodeWhispererClient(cfg, "")
	email := cwClient.FetchUserEmailFromAPI(ctx, accessToken, clientID, refreshToken)
	if email != "" {
		return email
	}

	// Method 2: Try SSO OIDC userinfo endpoint
	ssoClient := NewSSOOIDCClient(cfg)
	email = ssoClient.FetchUserEmail(ctx, accessToken)
	if email != "" {
		return email
	}

	// Method 3: Fallback to JWT parsing
	return ExtractEmailFromJWT(accessToken)
}