Skip to content

Commit ca99323

Browse files
luispaterluispater
authored
Merge pull request #42 from Ravens2121/master
feat(kiro): 新增授权码登录流程，优化邮箱获取与官方 Thinking 模式解析 预支持
2 parents 54c2fef + cf9a246 commit ca99323

File tree

9 files changed

+1126
-484
lines changed

9 files changed

+1126
-484
lines changed

‎cmd/server/main.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ func main() {
7878
var kiroLogin bool
7979
var kiroGoogleLogin bool
8080
var kiroAWSLogin bool
81+
var kiroAWSAuthCode bool
8182
var kiroImport bool
8283
var githubCopilotLogin bool
8384
var projectID string
@@ -101,6 +102,7 @@ func main() {
101102
flag.BoolVar(&kiroLogin, "kiro-login", false, "Login to Kiro using Google OAuth")
102103
flag.BoolVar(&kiroGoogleLogin, "kiro-google-login", false, "Login to Kiro using Google OAuth (same as --kiro-login)")
103104
flag.BoolVar(&kiroAWSLogin, "kiro-aws-login", false, "Login to Kiro using AWS Builder ID (device code flow)")
105+
flag.BoolVar(&kiroAWSAuthCode, "kiro-aws-authcode", false, "Login to Kiro using AWS Builder ID (authorization code flow, better UX)")
104106
flag.BoolVar(&kiroImport, "kiro-import", false, "Import Kiro token from Kiro IDE (~/.aws/sso/cache/kiro-auth-token.json)")
105107
flag.BoolVar(&githubCopilotLogin, "github-copilot-login", false, "Login to GitHub Copilot using device flow")
106108
flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")
@@ -513,6 +515,10 @@ func main() {
513515
// Users can explicitly override with --no-incognito
514516
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
515517
cmd.DoKiroAWSLogin(cfg, options)
518+
} else if kiroAWSAuthCode {
519+
// For Kiro auth with authorization code flow (better UX)
520+
setKiroIncognitoMode(cfg, useIncognito, noIncognito)
521+
cmd.DoKiroAWSAuthCodeLogin(cfg, options)
516522
} else if kiroImport {
517523
cmd.DoKiroImport(cfg, options)
518524
} else {

‎internal/auth/kiro/codewhisperer_client.go‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
// Package kiro provides CodeWhisperer API client for fetching user info.
2+
package kiro
3+
4+
import (
5+
"context"
6+
"encoding/json"
7+
"fmt"
8+
"io"
9+
"net/http"
10+
"time"
11+
12+
"github.com/google/uuid"
13+
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
14+
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
15+
log "github.com/sirupsen/logrus"
16+
)
17+
18+
const (
19+
codeWhispererAPI = "https://codewhisperer.us-east-1.amazonaws.com"
20+
kiroVersion = "0.6.18"
21+
)
22+
23+
// CodeWhispererClient handles CodeWhisperer API calls.
24+
type CodeWhispererClient struct {
25+
httpClient *http.Client
26+
machineID string
27+
}
28+
29+
// UsageLimitsResponse represents the getUsageLimits API response.
30+
type UsageLimitsResponse struct {
31+
DaysUntilReset *int `json:"daysUntilReset,omitempty"`
32+
NextDateReset *float64 `json:"nextDateReset,omitempty"`
33+
UserInfo *UserInfo `json:"userInfo,omitempty"`
34+
SubscriptionInfo *SubscriptionInfo `json:"subscriptionInfo,omitempty"`
35+
UsageBreakdownList []UsageBreakdown `json:"usageBreakdownList,omitempty"`
36+
}
37+
38+
// UserInfo contains user information from the API.
39+
type UserInfo struct {
40+
Email string `json:"email,omitempty"`
41+
UserID string `json:"userId,omitempty"`
42+
}
43+
44+
// SubscriptionInfo contains subscription details.
45+
type SubscriptionInfo struct {
46+
SubscriptionTitle string `json:"subscriptionTitle,omitempty"`
47+
Type string `json:"type,omitempty"`
48+
}
49+
50+
// UsageBreakdown contains usage details.
51+
type UsageBreakdown struct {
52+
UsageLimit *int `json:"usageLimit,omitempty"`
53+
CurrentUsage *int `json:"currentUsage,omitempty"`
54+
UsageLimitWithPrecision *float64 `json:"usageLimitWithPrecision,omitempty"`
55+
CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision,omitempty"`
56+
NextDateReset *float64 `json:"nextDateReset,omitempty"`
57+
DisplayName string `json:"displayName,omitempty"`
58+
ResourceType string `json:"resourceType,omitempty"`
59+
}
60+
61+
// NewCodeWhispererClient creates a new CodeWhisperer client.
62+
func NewCodeWhispererClient(cfg *config.Config, machineID string) *CodeWhispererClient {
63+
client := &http.Client{Timeout: 30 * time.Second}
64+
if cfg != nil {
65+
client = util.SetProxy(&cfg.SDKConfig, client)
66+
}
67+
if machineID == "" {
68+
machineID = uuid.New().String()
69+
}
70+
return &CodeWhispererClient{
71+
httpClient: client,
72+
machineID: machineID,
73+
}
74+
}
75+
76+
// generateInvocationID generates a unique invocation ID.
77+
func generateInvocationID() string {
78+
return uuid.New().String()
79+
}
80+
81+
// GetUsageLimits fetches usage limits and user info from CodeWhisperer API.
82+
// This is the recommended way to get user email after login.
83+
func (c *CodeWhispererClient) GetUsageLimits(ctx context.Context, accessToken string) (*UsageLimitsResponse, error) {
84+
url := fmt.Sprintf("%s/getUsageLimits?isEmailRequired=true&origin=AI_EDITOR&resourceType=AGENTIC_REQUEST", codeWhispererAPI)
85+
86+
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
87+
if err != nil {
88+
return nil, fmt.Errorf("failed to create request: %w", err)
89+
}
90+
91+
// Set headers to match Kiro IDE
92+
xAmzUserAgent := fmt.Sprintf("aws-sdk-js/1.0.0 KiroIDE-%s-%s", kiroVersion, c.machineID)
93+
userAgent := fmt.Sprintf("aws-sdk-js/1.0.0 ua/2.1 os/windows lang/js md/nodejs#20.16.0 api/codewhispererruntime#1.0.0 m/E KiroIDE-%s-%s", kiroVersion, c.machineID)
94+
95+
req.Header.Set("Authorization", "Bearer "+accessToken)
96+
req.Header.Set("x-amz-user-agent", xAmzUserAgent)
97+
req.Header.Set("User-Agent", userAgent)
98+
req.Header.Set("amz-sdk-invocation-id", generateInvocationID())
99+
req.Header.Set("amz-sdk-request", "attempt=1; max=1")
100+
req.Header.Set("Connection", "close")
101+
102+
log.Debugf("codewhisperer: GET %s", url)
103+
104+
resp, err := c.httpClient.Do(req)
105+
if err != nil {
106+
return nil, fmt.Errorf("request failed: %w", err)
107+
}
108+
defer resp.Body.Close()
109+
110+
body, err := io.ReadAll(resp.Body)
111+
if err != nil {
112+
return nil, fmt.Errorf("failed to read response: %w", err)
113+
}
114+
115+
log.Debugf("codewhisperer: status=%d, body=%s", resp.StatusCode, string(body))
116+
117+
if resp.StatusCode != http.StatusOK {
118+
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(body))
119+
}
120+
121+
var result UsageLimitsResponse
122+
if err := json.Unmarshal(body, &result); err != nil {
123+
return nil, fmt.Errorf("failed to parse response: %w", err)
124+
}
125+
126+
return &result, nil
127+
}
128+
129+
// FetchUserEmailFromAPI fetches user email using CodeWhisperer getUsageLimits API.
130+
// This is more reliable than JWT parsing as it uses the official API.
131+
func (c *CodeWhispererClient) FetchUserEmailFromAPI(ctx context.Context, accessToken string) string {
132+
resp, err := c.GetUsageLimits(ctx, accessToken)
133+
if err != nil {
134+
log.Debugf("codewhisperer: failed to get usage limits: %v", err)
135+
return ""
136+
}
137+
138+
if resp.UserInfo != nil && resp.UserInfo.Email != "" {
139+
log.Debugf("codewhisperer: got email from API: %s", resp.UserInfo.Email)
140+
return resp.UserInfo.Email
141+
}
142+
143+
log.Debugf("codewhisperer: no email in response")
144+
return ""
145+
}
146+
147+
// FetchUserEmailWithFallback fetches user email with multiple fallback methods.
148+
// Priority: 1. CodeWhisperer API 2. userinfo endpoint 3. JWT parsing
149+
func FetchUserEmailWithFallback(ctx context.Context, cfg *config.Config, accessToken string) string {
150+
// Method 1: Try CodeWhisperer API (most reliable)
151+
cwClient := NewCodeWhispererClient(cfg, "")
152+
email := cwClient.FetchUserEmailFromAPI(ctx, accessToken)
153+
if email != "" {
154+
return email
155+
}
156+
157+
// Method 2: Try SSO OIDC userinfo endpoint
158+
ssoClient := NewSSOOIDCClient(cfg)
159+
email = ssoClient.FetchUserEmail(ctx, accessToken)
160+
if email != "" {
161+
return email
162+
}
163+
164+
// Method 3: Fallback to JWT parsing
165+
return ExtractEmailFromJWT(accessToken)
166+
}

‎internal/auth/kiro/oauth.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,13 @@ func (o *KiroOAuth) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, err
163163
return ssoClient.LoginWithBuilderID(ctx)
164164
}
165165

166+
// LoginWithBuilderIDAuthCode performs OAuth login with AWS Builder ID using authorization code flow.
167+
// This provides a better UX than device code flow as it uses automatic browser callback.
168+
func (o *KiroOAuth) LoginWithBuilderIDAuthCode(ctx context.Context) (*KiroTokenData, error) {
169+
ssoClient := NewSSOOIDCClient(o.cfg)
170+
return ssoClient.LoginWithBuilderIDAuthCode(ctx)
171+
}
172+
166173
// exchangeCodeForToken exchanges the authorization code for tokens.
167174
func (o *KiroOAuth) exchangeCodeForToken(ctx context.Context, code, codeVerifier, redirectURI string) (*KiroTokenData, error) {
168175
payload := map[string]string{

0 commit comments

Comments
 (0)