feat(kiro): add AWS Builder ID authentication support

- Add --kiro-aws-login flag for AWS Builder ID device code flow
- Add DoKiroAWSLogin function for AWS SSO OIDC authentication
- Complete Kiro integration with AWS, Google OAuth, and social auth
- Add kiro executor, translator, and SDK components
- Update browser support for Kiro authentication flows
This commit is contained in:
Mansi
2025-12-05 22:46:24 +03:00
parent 0ebabf5152
commit 02d8a1cfec
45 changed files with 7519 additions and 171 deletions

2
.gitignore vendored
View File

@@ -1,5 +1,6 @@
# Binaries
cli-proxy-api
cliproxy
*.exe
# Configuration
@@ -31,6 +32,7 @@ GEMINI.md
.vscode/*
.claude/*
.serena/*
.mcp/cache/
# macOS
.DS_Store

View File

@@ -62,10 +62,16 @@ func main() {
var iflowCookie bool
var noBrowser bool
var antigravityLogin bool
var kiroLogin bool
var kiroGoogleLogin bool
var kiroAWSLogin bool
var kiroImport bool
var projectID string
var vertexImport string
var configPath string
var password string
var noIncognito bool
var useIncognito bool
// Define command-line flags for different operation modes.
flag.BoolVar(&login, "login", false, "Login Google Account")
@@ -75,7 +81,13 @@ func main() {
flag.BoolVar(&iflowLogin, "iflow-login", false, "Login to iFlow using OAuth")
flag.BoolVar(&iflowCookie, "iflow-cookie", false, "Login to iFlow using Cookie")
flag.BoolVar(&noBrowser, "no-browser", false, "Don't open browser automatically for OAuth")
flag.BoolVar(&useIncognito, "incognito", false, "Open browser in incognito/private mode for OAuth (useful for multiple accounts)")
flag.BoolVar(&noIncognito, "no-incognito", false, "Force disable incognito mode (uses existing browser session)")
flag.BoolVar(&antigravityLogin, "antigravity-login", false, "Login to Antigravity using OAuth")
flag.BoolVar(&kiroLogin, "kiro-login", false, "Login to Kiro using Google OAuth")
flag.BoolVar(&kiroGoogleLogin, "kiro-google-login", false, "Login to Kiro using Google OAuth (same as --kiro-login)")
flag.BoolVar(&kiroAWSLogin, "kiro-aws-login", false, "Login to Kiro using AWS Builder ID (device code flow)")
flag.BoolVar(&kiroImport, "kiro-import", false, "Import Kiro token from Kiro IDE (~/.aws/sso/cache/kiro-auth-token.json)")
flag.StringVar(&projectID, "project_id", "", "Project ID (Gemini only, not required)")
flag.StringVar(&configPath, "config", DefaultConfigPath, "Configure File Path")
flag.StringVar(&vertexImport, "vertex-import", "", "Import Vertex service account key JSON file")
@@ -448,6 +460,44 @@ func main() {
cmd.DoIFlowLogin(cfg, options)
} else if iflowCookie {
cmd.DoIFlowCookieAuth(cfg, options)
} else if kiroLogin {
// For Kiro auth, default to incognito mode for multi-account support
// Users can explicitly override with --no-incognito
// Note: This config mutation is safe - auth commands exit after completion
// and don't share config with StartService (which is in the else branch)
if useIncognito {
cfg.IncognitoBrowser = true
} else if noIncognito {
cfg.IncognitoBrowser = false
} else {
cfg.IncognitoBrowser = true // Kiro default
}
cmd.DoKiroLogin(cfg, options)
} else if kiroGoogleLogin {
// For Kiro auth, default to incognito mode for multi-account support
// Users can explicitly override with --no-incognito
// Note: This config mutation is safe - auth commands exit after completion
if useIncognito {
cfg.IncognitoBrowser = true
} else if noIncognito {
cfg.IncognitoBrowser = false
} else {
cfg.IncognitoBrowser = true // Kiro default
}
cmd.DoKiroGoogleLogin(cfg, options)
} else if kiroAWSLogin {
// For Kiro auth, default to incognito mode for multi-account support
// Users can explicitly override with --no-incognito
if useIncognito {
cfg.IncognitoBrowser = true
} else if noIncognito {
cfg.IncognitoBrowser = false
} else {
cfg.IncognitoBrowser = true // Kiro default
}
cmd.DoKiroAWSLogin(cfg, options)
} else if kiroImport {
cmd.DoKiroImport(cfg, options)
} else {
// In cloud deploy mode without config file, just wait for shutdown signals
if isCloudDeploy && !configFileExists {

View File

@@ -32,6 +32,11 @@ api-keys:
# Enable debug logging
debug: false
# Open OAuth URLs in incognito/private browser mode.
# Useful when you want to login with a different account without logging out from your current session.
# Default: false (but Kiro auth defaults to true for multi-account support)
incognito-browser: true
# When true, write application logs to rotating files instead of stdout
logging-to-file: false
@@ -99,6 +104,16 @@ ws-auth: false
# - "*-think" # wildcard matching suffix (e.g. claude-opus-4-5-thinking)
# - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022)
# Kiro (AWS CodeWhisperer) configuration
# Note: Kiro API currently only operates in us-east-1 region
#kiro:
# - token-file: "~/.aws/sso/cache/kiro-auth-token.json" # path to Kiro token file
# agent-task-type: "" # optional: "vibe" or empty (API default)
# - access-token: "aoaAAAAA..." # or provide tokens directly
# refresh-token: "aorAAAAA..."
# profile-arn: "arn:aws:codewhisperer:us-east-1:..."
# proxy-url: "socks5://proxy.example.com:1080" # optional: proxy override
# OpenAI compatibility providers
# openai-compatibility:
# - name: "openrouter" # The name of the provider; it will be used in the user agent and other places.

3
go.mod
View File

@@ -13,14 +13,15 @@ require (
github.com/joho/godotenv v1.5.1
github.com/klauspost/compress v1.17.4
github.com/minio/minio-go/v7 v7.0.66
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
github.com/sirupsen/logrus v1.9.3
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
github.com/tidwall/gjson v1.18.0
github.com/tidwall/sjson v1.2.5
github.com/tiktoken-go/tokenizer v0.7.0
golang.org/x/crypto v0.43.0
golang.org/x/net v0.46.0
golang.org/x/oauth2 v0.30.0
golang.org/x/term v0.36.0
gopkg.in/natefinch/lumberjack.v2 v2.2.1
gopkg.in/yaml.v3 v3.0.1
)

5
go.sum
View File

@@ -116,6 +116,8 @@ github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6
github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
github.com/pjbgf/sha1cd v0.5.0 h1:a+UkboSi1znleCDUNT3M5YxjOnN1fz2FhN48FlwCxs0=
github.com/pjbgf/sha1cd v0.5.0/go.mod h1:lhpGlyHLpQZoxMv8HcgXvZEhcGs0PG/vsZnEJ7H0iCM=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
@@ -126,8 +128,6 @@ github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw=
github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA=
github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966/go.mod h1:sUM3LWHvSMaG192sy56D9F7CNvL7jUJVXoqM1QKLnog=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
@@ -169,6 +169,7 @@ golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKl
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=

View File

@@ -36,9 +36,32 @@ import (
)
var (
oauthStatus = make(map[string]string)
oauthStatus = make(map[string]string)
oauthStatusMutex sync.RWMutex
)
// getOAuthStatus safely retrieves an OAuth status
func getOAuthStatus(key string) (string, bool) {
oauthStatusMutex.RLock()
defer oauthStatusMutex.RUnlock()
status, ok := oauthStatus[key]
return status, ok
}
// setOAuthStatus safely sets an OAuth status
func setOAuthStatus(key string, status string) {
oauthStatusMutex.Lock()
defer oauthStatusMutex.Unlock()
oauthStatus[key] = status
}
// deleteOAuthStatus safely deletes an OAuth status
func deleteOAuthStatus(key string) {
oauthStatusMutex.Lock()
defer oauthStatusMutex.Unlock()
delete(oauthStatus, key)
}
var lastRefreshKeys = []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"}
const (
@@ -760,7 +783,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
deadline := time.Now().Add(timeout)
for {
if time.Now().After(deadline) {
oauthStatus[state] = "Timeout waiting for OAuth callback"
setOAuthStatus(state, "Timeout waiting for OAuth callback")
return nil, fmt.Errorf("timeout waiting for OAuth callback")
}
data, errRead := os.ReadFile(path)
@@ -785,13 +808,13 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
if errStr := resultMap["error"]; errStr != "" {
oauthErr := claude.NewOAuthError(errStr, "", http.StatusBadRequest)
log.Error(claude.GetUserFriendlyMessage(oauthErr))
oauthStatus[state] = "Bad request"
setOAuthStatus(state, "Bad request")
return
}
if resultMap["state"] != state {
authErr := claude.NewAuthenticationError(claude.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, resultMap["state"]))
log.Error(claude.GetUserFriendlyMessage(authErr))
oauthStatus[state] = "State code error"
setOAuthStatus(state, "State code error")
return
}
@@ -824,7 +847,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
if errDo != nil {
authErr := claude.NewAuthenticationError(claude.ErrCodeExchangeFailed, errDo)
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
oauthStatus[state] = "Failed to exchange authorization code for tokens"
setOAuthStatus(state, "Failed to exchange authorization code for tokens")
return
}
defer func() {
@@ -835,7 +858,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
log.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(respBody))
oauthStatus[state] = fmt.Sprintf("token exchange failed with status %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("token exchange failed with status %d", resp.StatusCode))
return
}
var tResp struct {
@@ -848,7 +871,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
}
if errU := json.Unmarshal(respBody, &tResp); errU != nil {
log.Errorf("failed to parse token response: %v", errU)
oauthStatus[state] = "Failed to parse token response"
setOAuthStatus(state, "Failed to parse token response")
return
}
bundle := &claude.ClaudeAuthBundle{
@@ -873,7 +896,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save authentication tokens: %v", errSave)
oauthStatus[state] = "Failed to save authentication tokens"
setOAuthStatus(state, "Failed to save authentication tokens")
return
}
@@ -882,10 +905,10 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
fmt.Println("API key obtained and saved")
}
fmt.Println("You can now use Claude services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -944,7 +967,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
for {
if time.Now().After(deadline) {
log.Error("oauth flow timed out")
oauthStatus[state] = "OAuth flow timed out"
setOAuthStatus(state, "OAuth flow timed out")
return
}
if data, errR := os.ReadFile(waitFile); errR == nil {
@@ -953,13 +976,13 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
_ = os.Remove(waitFile)
if errStr := m["error"]; errStr != "" {
log.Errorf("Authentication failed: %s", errStr)
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
return
}
authCode = m["code"]
if authCode == "" {
log.Errorf("Authentication failed: code not found")
oauthStatus[state] = "Authentication failed: code not found"
setOAuthStatus(state, "Authentication failed: code not found")
return
}
break
@@ -971,7 +994,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
token, err := conf.Exchange(ctx, authCode)
if err != nil {
log.Errorf("Failed to exchange token: %v", err)
oauthStatus[state] = "Failed to exchange token"
setOAuthStatus(state, "Failed to exchange token")
return
}
@@ -982,7 +1005,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
req, errNewRequest := http.NewRequestWithContext(ctx, "GET", "https://www.googleapis.com/oauth2/v1/userinfo?alt=json", nil)
if errNewRequest != nil {
log.Errorf("Could not get user info: %v", errNewRequest)
oauthStatus[state] = "Could not get user info"
setOAuthStatus(state, "Could not get user info")
return
}
req.Header.Set("Content-Type", "application/json")
@@ -991,7 +1014,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
resp, errDo := authHTTPClient.Do(req)
if errDo != nil {
log.Errorf("Failed to execute request: %v", errDo)
oauthStatus[state] = "Failed to execute request"
setOAuthStatus(state, "Failed to execute request")
return
}
defer func() {
@@ -1003,7 +1026,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
bodyBytes, _ := io.ReadAll(resp.Body)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
log.Errorf("Get user info request failed with status %d: %s", resp.StatusCode, string(bodyBytes))
oauthStatus[state] = fmt.Sprintf("Get user info request failed with status %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("Get user info request failed with status %d", resp.StatusCode))
return
}
@@ -1012,7 +1035,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
fmt.Printf("Authenticated user email: %s\n", email)
} else {
fmt.Println("Failed to get user email from token")
oauthStatus[state] = "Failed to get user email from token"
setOAuthStatus(state, "Failed to get user email from token")
}
// Marshal/unmarshal oauth2.Token to generic map and enrich fields
@@ -1020,7 +1043,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
jsonData, _ := json.Marshal(token)
if errUnmarshal := json.Unmarshal(jsonData, &ifToken); errUnmarshal != nil {
log.Errorf("Failed to unmarshal token: %v", errUnmarshal)
oauthStatus[state] = "Failed to unmarshal token"
setOAuthStatus(state, "Failed to unmarshal token")
return
}
@@ -1046,7 +1069,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
gemClient, errGetClient := gemAuth.GetAuthenticatedClient(ctx, &ts, h.cfg, true)
if errGetClient != nil {
log.Fatalf("failed to get authenticated client: %v", errGetClient)
oauthStatus[state] = "Failed to get authenticated client"
setOAuthStatus(state, "Failed to get authenticated client")
return
}
fmt.Println("Authentication successful.")
@@ -1056,12 +1079,12 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
projects, errAll := onboardAllGeminiProjects(ctx, gemClient, &ts)
if errAll != nil {
log.Errorf("Failed to complete Gemini CLI onboarding: %v", errAll)
oauthStatus[state] = "Failed to complete Gemini CLI onboarding"
setOAuthStatus(state, "Failed to complete Gemini CLI onboarding")
return
}
if errVerify := ensureGeminiProjectsEnabled(ctx, gemClient, projects); errVerify != nil {
log.Errorf("Failed to verify Cloud AI API status: %v", errVerify)
oauthStatus[state] = "Failed to verify Cloud AI API status"
setOAuthStatus(state, "Failed to verify Cloud AI API status")
return
}
ts.ProjectID = strings.Join(projects, ",")
@@ -1069,26 +1092,26 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
} else {
if errEnsure := ensureGeminiProjectAndOnboard(ctx, gemClient, &ts, requestedProjectID); errEnsure != nil {
log.Errorf("Failed to complete Gemini CLI onboarding: %v", errEnsure)
oauthStatus[state] = "Failed to complete Gemini CLI onboarding"
setOAuthStatus(state, "Failed to complete Gemini CLI onboarding")
return
}
if strings.TrimSpace(ts.ProjectID) == "" {
log.Error("Onboarding did not return a project ID")
oauthStatus[state] = "Failed to resolve project ID"
setOAuthStatus(state, "Failed to resolve project ID")
return
}
isChecked, errCheck := checkCloudAPIIsEnabled(ctx, gemClient, ts.ProjectID)
if errCheck != nil {
log.Errorf("Failed to verify Cloud AI API status: %v", errCheck)
oauthStatus[state] = "Failed to verify Cloud AI API status"
setOAuthStatus(state, "Failed to verify Cloud AI API status")
return
}
ts.Checked = isChecked
if !isChecked {
log.Error("Cloud AI API is not enabled for the selected project")
oauthStatus[state] = "Cloud AI API not enabled"
setOAuthStatus(state, "Cloud AI API not enabled")
return
}
}
@@ -1111,15 +1134,15 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save token to file: %v", errSave)
oauthStatus[state] = "Failed to save token to file"
setOAuthStatus(state, "Failed to save token to file")
return
}
delete(oauthStatus, state)
deleteOAuthStatus(state)
fmt.Printf("You can now use Gemini CLI services through this CLI; token saved to %s\n", savedPath)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1180,7 +1203,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
if time.Now().After(deadline) {
authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, fmt.Errorf("timeout waiting for OAuth callback"))
log.Error(codex.GetUserFriendlyMessage(authErr))
oauthStatus[state] = "Timeout waiting for OAuth callback"
setOAuthStatus(state, "Timeout waiting for OAuth callback")
return
}
if data, errR := os.ReadFile(waitFile); errR == nil {
@@ -1190,12 +1213,12 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
if errStr := m["error"]; errStr != "" {
oauthErr := codex.NewOAuthError(errStr, "", http.StatusBadRequest)
log.Error(codex.GetUserFriendlyMessage(oauthErr))
oauthStatus[state] = "Bad Request"
setOAuthStatus(state, "Bad Request")
return
}
if m["state"] != state {
authErr := codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, m["state"]))
oauthStatus[state] = "State code error"
setOAuthStatus(state, "State code error")
log.Error(codex.GetUserFriendlyMessage(authErr))
return
}
@@ -1226,14 +1249,14 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
resp, errDo := httpClient.Do(req)
if errDo != nil {
authErr := codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, errDo)
oauthStatus[state] = "Failed to exchange authorization code for tokens"
setOAuthStatus(state, "Failed to exchange authorization code for tokens")
log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
return
}
defer func() { _ = resp.Body.Close() }()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
oauthStatus[state] = fmt.Sprintf("Token exchange failed with status %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("Token exchange failed with status %d", resp.StatusCode))
log.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(respBody))
return
}
@@ -1244,7 +1267,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
ExpiresIn int `json:"expires_in"`
}
if errU := json.Unmarshal(respBody, &tokenResp); errU != nil {
oauthStatus[state] = "Failed to parse token response"
setOAuthStatus(state, "Failed to parse token response")
log.Errorf("failed to parse token response: %v", errU)
return
}
@@ -1282,7 +1305,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
}
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
oauthStatus[state] = "Failed to save authentication tokens"
setOAuthStatus(state, "Failed to save authentication tokens")
log.Fatalf("Failed to save authentication tokens: %v", errSave)
return
}
@@ -1291,10 +1314,10 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
fmt.Println("API key obtained and saved")
}
fmt.Println("You can now use Codex services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1360,7 +1383,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
for {
if time.Now().After(deadline) {
log.Error("oauth flow timed out")
oauthStatus[state] = "OAuth flow timed out"
setOAuthStatus(state, "OAuth flow timed out")
return
}
if data, errReadFile := os.ReadFile(waitFile); errReadFile == nil {
@@ -1369,18 +1392,18 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
_ = os.Remove(waitFile)
if errStr := strings.TrimSpace(payload["error"]); errStr != "" {
log.Errorf("Authentication failed: %s", errStr)
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
return
}
if payloadState := strings.TrimSpace(payload["state"]); payloadState != "" && payloadState != state {
log.Errorf("Authentication failed: state mismatch")
oauthStatus[state] = "Authentication failed: state mismatch"
setOAuthStatus(state, "Authentication failed: state mismatch")
return
}
authCode = strings.TrimSpace(payload["code"])
if authCode == "" {
log.Error("Authentication failed: code not found")
oauthStatus[state] = "Authentication failed: code not found"
setOAuthStatus(state, "Authentication failed: code not found")
return
}
break
@@ -1399,7 +1422,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
req, errNewRequest := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode()))
if errNewRequest != nil {
log.Errorf("Failed to build token request: %v", errNewRequest)
oauthStatus[state] = "Failed to build token request"
setOAuthStatus(state, "Failed to build token request")
return
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
@@ -1407,7 +1430,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
resp, errDo := httpClient.Do(req)
if errDo != nil {
log.Errorf("Failed to execute token request: %v", errDo)
oauthStatus[state] = "Failed to exchange token"
setOAuthStatus(state, "Failed to exchange token")
return
}
defer func() {
@@ -1419,7 +1442,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
bodyBytes, _ := io.ReadAll(resp.Body)
log.Errorf("Antigravity token exchange failed with status %d: %s", resp.StatusCode, string(bodyBytes))
oauthStatus[state] = fmt.Sprintf("Token exchange failed: %d", resp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("Token exchange failed: %d", resp.StatusCode))
return
}
@@ -1431,7 +1454,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
}
if errDecode := json.NewDecoder(resp.Body).Decode(&tokenResp); errDecode != nil {
log.Errorf("Failed to parse token response: %v", errDecode)
oauthStatus[state] = "Failed to parse token response"
setOAuthStatus(state, "Failed to parse token response")
return
}
@@ -1440,7 +1463,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
infoReq, errInfoReq := http.NewRequestWithContext(ctx, http.MethodGet, "https://www.googleapis.com/oauth2/v1/userinfo?alt=json", nil)
if errInfoReq != nil {
log.Errorf("Failed to build user info request: %v", errInfoReq)
oauthStatus[state] = "Failed to build user info request"
setOAuthStatus(state, "Failed to build user info request")
return
}
infoReq.Header.Set("Authorization", "Bearer "+tokenResp.AccessToken)
@@ -1448,7 +1471,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
infoResp, errInfo := httpClient.Do(infoReq)
if errInfo != nil {
log.Errorf("Failed to execute user info request: %v", errInfo)
oauthStatus[state] = "Failed to execute user info request"
setOAuthStatus(state, "Failed to execute user info request")
return
}
defer func() {
@@ -1467,7 +1490,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
} else {
bodyBytes, _ := io.ReadAll(infoResp.Body)
log.Errorf("User info request failed with status %d: %s", infoResp.StatusCode, string(bodyBytes))
oauthStatus[state] = fmt.Sprintf("User info request failed: %d", infoResp.StatusCode)
setOAuthStatus(state, fmt.Sprintf("User info request failed: %d", infoResp.StatusCode))
return
}
}
@@ -1515,11 +1538,11 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save token to file: %v", errSave)
oauthStatus[state] = "Failed to save token to file"
setOAuthStatus(state, "Failed to save token to file")
return
}
delete(oauthStatus, state)
deleteOAuthStatus(state)
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
if projectID != "" {
fmt.Printf("Using GCP project: %s\n", projectID)
@@ -1527,7 +1550,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
fmt.Println("You can now use Antigravity services through this CLI")
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1552,7 +1575,7 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
fmt.Println("Waiting for authentication...")
tokenData, errPollForToken := qwenAuth.PollForToken(deviceFlow.DeviceCode, deviceFlow.CodeVerifier)
if errPollForToken != nil {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Printf("Authentication failed: %v\n", errPollForToken)
return
}
@@ -1571,16 +1594,16 @@ func (h *Handler) RequestQwenToken(c *gin.Context) {
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
log.Fatalf("Failed to save authentication tokens: %v", errSave)
oauthStatus[state] = "Failed to save authentication tokens"
setOAuthStatus(state, "Failed to save authentication tokens")
return
}
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
fmt.Println("You can now use Qwen services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -1619,7 +1642,7 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
var resultMap map[string]string
for {
if time.Now().After(deadline) {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Println("Authentication failed: timeout waiting for callback")
return
}
@@ -1632,26 +1655,26 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
}
if errStr := strings.TrimSpace(resultMap["error"]); errStr != "" {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Printf("Authentication failed: %s\n", errStr)
return
}
if resultState := strings.TrimSpace(resultMap["state"]); resultState != state {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Println("Authentication failed: state mismatch")
return
}
code := strings.TrimSpace(resultMap["code"])
if code == "" {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Println("Authentication failed: code missing")
return
}
tokenData, errExchange := authSvc.ExchangeCodeForTokens(ctx, code, redirectURI)
if errExchange != nil {
oauthStatus[state] = "Authentication failed"
setOAuthStatus(state, "Authentication failed")
fmt.Printf("Authentication failed: %v\n", errExchange)
return
}
@@ -1673,7 +1696,7 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
savedPath, errSave := h.saveTokenRecord(ctx, record)
if errSave != nil {
oauthStatus[state] = "Failed to save authentication tokens"
setOAuthStatus(state, "Failed to save authentication tokens")
log.Fatalf("Failed to save authentication tokens: %v", errSave)
return
}
@@ -1683,10 +1706,10 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
fmt.Println("API key obtained and saved")
}
fmt.Println("You can now use iFlow services through this CLI")
delete(oauthStatus, state)
deleteOAuthStatus(state)
}()
oauthStatus[state] = ""
setOAuthStatus(state, "")
c.JSON(http.StatusOK, gin.H{"status": "ok", "url": authURL, "state": state})
}
@@ -2110,7 +2133,7 @@ func checkCloudAPIIsEnabled(ctx context.Context, httpClient *http.Client, projec
func (h *Handler) GetAuthStatus(c *gin.Context) {
state := c.Query("state")
if err, ok := oauthStatus[state]; ok {
if err, ok := getOAuthStatus(state); ok {
if err != "" {
c.JSON(200, gin.H{"status": "error", "error": err})
} else {
@@ -2120,5 +2143,5 @@ func (h *Handler) GetAuthStatus(c *gin.Context) {
} else {
c.JSON(200, gin.H{"status": "ok"})
}
delete(oauthStatus, state)
deleteOAuthStatus(state)
}

View File

@@ -902,7 +902,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
for _, p := range cfg.OpenAICompatibility {
providerNames = append(providerNames, p.Name)
}
s.handlers.OpenAICompatProviders = providerNames
s.handlers.SetOpenAICompatProviders(providerNames)
s.handlers.UpdateClients(&cfg.SDKConfig)

View File

@@ -242,6 +242,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
platformURL = "https://console.anthropic.com/"
}
// Validate platformURL to prevent XSS - only allow http/https URLs
if !isValidURL(platformURL) {
platformURL = "https://console.anthropic.com/"
}
// Generate success page HTML with dynamic content
successHTML := s.generateSuccessHTML(setupRequired, platformURL)
@@ -251,6 +256,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
}
}
// isValidURL checks if the URL is a valid http/https URL to prevent XSS
func isValidURL(urlStr string) bool {
urlStr = strings.TrimSpace(urlStr)
return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
}
// generateSuccessHTML creates the HTML content for the success page.
// It customizes the page based on whether additional setup is required
// and includes a link to the platform.

View File

@@ -239,6 +239,11 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
platformURL = "https://platform.openai.com"
}
// Validate platformURL to prevent XSS - only allow http/https URLs
if !isValidURL(platformURL) {
platformURL = "https://platform.openai.com"
}
// Generate success page HTML with dynamic content
successHTML := s.generateSuccessHTML(setupRequired, platformURL)
@@ -248,6 +253,12 @@ func (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {
}
}
// isValidURL checks if the URL is a valid http/https URL to prevent XSS
func isValidURL(urlStr string) bool {
urlStr = strings.TrimSpace(urlStr)
return strings.HasPrefix(urlStr, "https://") || strings.HasPrefix(urlStr, "http://")
}
// generateSuccessHTML creates the HTML content for the success page.
// It customizes the page based on whether additional setup is required
// and includes a link to the platform.

View File

@@ -9,6 +9,7 @@ import (
"io"
"net/http"
"net/url"
"os"
"strings"
"time"
@@ -28,10 +29,21 @@ const (
iFlowAPIKeyEndpoint = "https://platform.iflow.cn/api/openapi/apikey"
// Client credentials provided by iFlow for the Code Assist integration.
iFlowOAuthClientID = "10009311001"
iFlowOAuthClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"
iFlowOAuthClientID = "10009311001"
// Default client secret (can be overridden via IFLOW_CLIENT_SECRET env var)
defaultIFlowClientSecret = "4Z3YjXycVsQvyGF1etiNlIBB4RsqSDtW"
)
// getIFlowClientSecret returns the iFlow OAuth client secret.
// It first checks the IFLOW_CLIENT_SECRET environment variable,
// falling back to the default value if not set.
func getIFlowClientSecret() string {
if secret := os.Getenv("IFLOW_CLIENT_SECRET"); secret != "" {
return secret
}
return defaultIFlowClientSecret
}
// DefaultAPIBaseURL is the canonical chat completions endpoint.
const DefaultAPIBaseURL = "https://apis.iflow.cn/v1"
@@ -72,7 +84,7 @@ func (ia *IFlowAuth) ExchangeCodeForTokens(ctx context.Context, code, redirectUR
form.Set("code", code)
form.Set("redirect_uri", redirectURI)
form.Set("client_id", iFlowOAuthClientID)
form.Set("client_secret", iFlowOAuthClientSecret)
form.Set("client_secret", getIFlowClientSecret())
req, err := ia.newTokenRequest(ctx, form)
if err != nil {
@@ -88,7 +100,7 @@ func (ia *IFlowAuth) RefreshTokens(ctx context.Context, refreshToken string) (*I
form.Set("grant_type", "refresh_token")
form.Set("refresh_token", refreshToken)
form.Set("client_id", iFlowOAuthClientID)
form.Set("client_secret", iFlowOAuthClientSecret)
form.Set("client_secret", getIFlowClientSecret())
req, err := ia.newTokenRequest(ctx, form)
if err != nil {
@@ -104,7 +116,7 @@ func (ia *IFlowAuth) newTokenRequest(ctx context.Context, form url.Values) (*htt
return nil, fmt.Errorf("iflow token: create request failed: %w", err)
}
basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + iFlowOAuthClientSecret))
basic := base64.StdEncoding.EncodeToString([]byte(iFlowOAuthClientID + ":" + getIFlowClientSecret()))
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Basic "+basic)

301
internal/auth/kiro/aws.go Normal file
View File

@@ -0,0 +1,301 @@
// Package kiro provides authentication functionality for AWS CodeWhisperer (Kiro) API.
// It includes interfaces and implementations for token storage and authentication methods.
package kiro
import (
"encoding/base64"
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
)
// PKCECodes holds PKCE verification codes for OAuth2 PKCE flow
type PKCECodes struct {
// CodeVerifier is the cryptographically random string used to correlate
// the authorization request to the token request
CodeVerifier string `json:"code_verifier"`
// CodeChallenge is the SHA256 hash of the code verifier, base64url-encoded
CodeChallenge string `json:"code_challenge"`
}
// KiroTokenData holds OAuth token information from AWS CodeWhisperer (Kiro)
type KiroTokenData struct {
// AccessToken is the OAuth2 access token for API access
AccessToken string `json:"accessToken"`
// RefreshToken is used to obtain new access tokens
RefreshToken string `json:"refreshToken"`
// ProfileArn is the AWS CodeWhisperer profile ARN
ProfileArn string `json:"profileArn"`
// ExpiresAt is the timestamp when the token expires
ExpiresAt string `json:"expiresAt"`
// AuthMethod indicates the authentication method used (e.g., "builder-id", "social")
AuthMethod string `json:"authMethod"`
// Provider indicates the OAuth provider (e.g., "AWS", "Google")
Provider string `json:"provider"`
// ClientID is the OIDC client ID (needed for token refresh)
ClientID string `json:"clientId,omitempty"`
// ClientSecret is the OIDC client secret (needed for token refresh)
ClientSecret string `json:"clientSecret,omitempty"`
// Email is the user's email address (used for file naming)
Email string `json:"email,omitempty"`
}
// KiroAuthBundle aggregates authentication data after OAuth flow completion
type KiroAuthBundle struct {
// TokenData contains the OAuth tokens from the authentication flow
TokenData KiroTokenData `json:"token_data"`
// LastRefresh is the timestamp of the last token refresh
LastRefresh string `json:"last_refresh"`
}
// KiroUsageInfo represents usage information from CodeWhisperer API
type KiroUsageInfo struct {
// SubscriptionTitle is the subscription plan name (e.g., "KIRO FREE")
SubscriptionTitle string `json:"subscription_title"`
// CurrentUsage is the current credit usage
CurrentUsage float64 `json:"current_usage"`
// UsageLimit is the maximum credit limit
UsageLimit float64 `json:"usage_limit"`
// NextReset is the timestamp of the next usage reset
NextReset string `json:"next_reset"`
}
// KiroModel represents a model available through the CodeWhisperer API
type KiroModel struct {
// ModelID is the unique identifier for the model
ModelID string `json:"modelId"`
// ModelName is the human-readable name
ModelName string `json:"modelName"`
// Description is the model description
Description string `json:"description"`
// RateMultiplier is the credit multiplier for this model
RateMultiplier float64 `json:"rateMultiplier"`
// RateUnit is the unit for rate calculation (e.g., "credit")
RateUnit string `json:"rateUnit"`
// MaxInputTokens is the maximum input token limit
MaxInputTokens int `json:"maxInputTokens,omitempty"`
}
// KiroIDETokenFile is the default path to Kiro IDE's token file
const KiroIDETokenFile = ".aws/sso/cache/kiro-auth-token.json"
// LoadKiroIDEToken loads token data from Kiro IDE's token file.
func LoadKiroIDEToken() (*KiroTokenData, error) {
homeDir, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
tokenPath := filepath.Join(homeDir, KiroIDETokenFile)
data, err := os.ReadFile(tokenPath)
if err != nil {
return nil, fmt.Errorf("failed to read Kiro IDE token file (%s): %w", tokenPath, err)
}
var token KiroTokenData
if err := json.Unmarshal(data, &token); err != nil {
return nil, fmt.Errorf("failed to parse Kiro IDE token: %w", err)
}
if token.AccessToken == "" {
return nil, fmt.Errorf("access token is empty in Kiro IDE token file")
}
return &token, nil
}
// LoadKiroTokenFromPath loads token data from a custom path.
// This supports multiple accounts by allowing different token files.
func LoadKiroTokenFromPath(tokenPath string) (*KiroTokenData, error) {
// Expand ~ to home directory
if len(tokenPath) > 0 && tokenPath[0] == '~' {
homeDir, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
tokenPath = filepath.Join(homeDir, tokenPath[1:])
}
data, err := os.ReadFile(tokenPath)
if err != nil {
return nil, fmt.Errorf("failed to read token file (%s): %w", tokenPath, err)
}
var token KiroTokenData
if err := json.Unmarshal(data, &token); err != nil {
return nil, fmt.Errorf("failed to parse token file: %w", err)
}
if token.AccessToken == "" {
return nil, fmt.Errorf("access token is empty in token file")
}
return &token, nil
}
// ListKiroTokenFiles lists all Kiro token files in the cache directory.
// This supports multiple accounts by finding all token files.
func ListKiroTokenFiles() ([]string, error) {
homeDir, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
cacheDir := filepath.Join(homeDir, ".aws", "sso", "cache")
// Check if directory exists
if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
return nil, nil // No token files
}
entries, err := os.ReadDir(cacheDir)
if err != nil {
return nil, fmt.Errorf("failed to read cache directory: %w", err)
}
var tokenFiles []string
for _, entry := range entries {
if entry.IsDir() {
continue
}
name := entry.Name()
// Look for kiro token files only (avoid matching unrelated AWS SSO cache files)
if strings.HasSuffix(name, ".json") && strings.HasPrefix(name, "kiro") {
tokenFiles = append(tokenFiles, filepath.Join(cacheDir, name))
}
}
return tokenFiles, nil
}
// LoadAllKiroTokens loads all Kiro tokens from the cache directory.
// This supports multiple accounts.
func LoadAllKiroTokens() ([]*KiroTokenData, error) {
files, err := ListKiroTokenFiles()
if err != nil {
return nil, err
}
var tokens []*KiroTokenData
for _, file := range files {
token, err := LoadKiroTokenFromPath(file)
if err != nil {
// Skip invalid token files
continue
}
tokens = append(tokens, token)
}
return tokens, nil
}
// JWTClaims represents the claims we care about from a JWT token.
// JWT tokens from Kiro/AWS contain user information in the payload.
type JWTClaims struct {
Email string `json:"email,omitempty"`
Sub string `json:"sub,omitempty"`
PreferredUser string `json:"preferred_username,omitempty"`
Name string `json:"name,omitempty"`
Iss string `json:"iss,omitempty"`
}
// ExtractEmailFromJWT extracts the user's email from a JWT access token.
// JWT tokens typically have format: header.payload.signature
// The payload is base64url-encoded JSON containing user claims.
func ExtractEmailFromJWT(accessToken string) string {
if accessToken == "" {
return ""
}
// JWT format: header.payload.signature
parts := strings.Split(accessToken, ".")
if len(parts) != 3 {
return ""
}
// Decode the payload (second part)
payload := parts[1]
// Add padding if needed (base64url requires padding)
switch len(payload) % 4 {
case 2:
payload += "=="
case 3:
payload += "="
}
decoded, err := base64.URLEncoding.DecodeString(payload)
if err != nil {
// Try RawURLEncoding (no padding)
decoded, err = base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return ""
}
}
var claims JWTClaims
if err := json.Unmarshal(decoded, &claims); err != nil {
return ""
}
// Return email if available
if claims.Email != "" {
return claims.Email
}
// Fallback to preferred_username (some providers use this)
if claims.PreferredUser != "" && strings.Contains(claims.PreferredUser, "@") {
return claims.PreferredUser
}
// Fallback to sub if it looks like an email
if claims.Sub != "" && strings.Contains(claims.Sub, "@") {
return claims.Sub
}
return ""
}
// SanitizeEmailForFilename sanitizes an email address for use in a filename.
// Replaces special characters with underscores and prevents path traversal attacks.
// Also handles URL-encoded characters to prevent encoded path traversal attempts.
func SanitizeEmailForFilename(email string) string {
if email == "" {
return ""
}
result := email
// First, handle URL-encoded path traversal attempts (%2F, %2E, %5C, etc.)
// This prevents encoded characters from bypassing the sanitization.
// Note: We replace % last to catch any remaining encodings including double-encoding (%252F)
result = strings.ReplaceAll(result, "%2F", "_") // /
result = strings.ReplaceAll(result, "%2f", "_")
result = strings.ReplaceAll(result, "%5C", "_") // \
result = strings.ReplaceAll(result, "%5c", "_")
result = strings.ReplaceAll(result, "%2E", "_") // .
result = strings.ReplaceAll(result, "%2e", "_")
result = strings.ReplaceAll(result, "%00", "_") // null byte
result = strings.ReplaceAll(result, "%", "_") // Catch remaining % to prevent double-encoding attacks
// Replace characters that are problematic in filenames
// Keep @ and . in middle but replace other special characters
for _, char := range []string{"/", "\\", ":", "*", "?", "\"", "<", ">", "|", " ", "\x00"} {
result = strings.ReplaceAll(result, char, "_")
}
// Prevent path traversal: replace leading dots in each path component
// This handles cases like "../../../etc/passwd" → "_.._.._.._etc_passwd"
parts := strings.Split(result, "_")
for i, part := range parts {
for strings.HasPrefix(part, ".") {
part = "_" + part[1:]
}
parts[i] = part
}
result = strings.Join(parts, "_")
return result
}

View File

@@ -0,0 +1,314 @@
// Package kiro provides OAuth2 authentication functionality for AWS CodeWhisperer (Kiro) API.
// This package implements token loading, refresh, and API communication with CodeWhisperer.
package kiro
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
const (
// awsKiroEndpoint is used for CodeWhisperer management APIs (GetUsageLimits, ListProfiles, etc.)
// Note: This is different from the Amazon Q streaming endpoint (q.us-east-1.amazonaws.com)
// used in kiro_executor.go for GenerateAssistantResponse. Both endpoints are correct
// for their respective API operations.
awsKiroEndpoint = "https://codewhisperer.us-east-1.amazonaws.com"
defaultTokenFile = "~/.aws/sso/cache/kiro-auth-token.json"
targetGetUsage = "AmazonCodeWhispererService.GetUsageLimits"
targetListModels = "AmazonCodeWhispererService.ListAvailableModels"
targetGenerateChat = "AmazonCodeWhispererStreamingService.GenerateAssistantResponse"
)
// KiroAuth handles AWS CodeWhisperer authentication and API communication.
// It provides methods for loading tokens, refreshing expired tokens,
// and communicating with the CodeWhisperer API.
type KiroAuth struct {
httpClient *http.Client
endpoint string
}
// NewKiroAuth creates a new Kiro authentication service.
// It initializes the HTTP client with proxy settings from the configuration.
//
// Parameters:
// - cfg: The application configuration containing proxy settings
//
// Returns:
// - *KiroAuth: A new Kiro authentication service instance
func NewKiroAuth(cfg *config.Config) *KiroAuth {
return &KiroAuth{
httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{Timeout: 120 * time.Second}),
endpoint: awsKiroEndpoint,
}
}
// LoadTokenFromFile loads token data from a file path.
// This method reads and parses the token file, expanding ~ to the home directory.
//
// Parameters:
// - tokenFile: Path to the token file (supports ~ expansion)
//
// Returns:
// - *KiroTokenData: The parsed token data
// - error: An error if file reading or parsing fails
func (k *KiroAuth) LoadTokenFromFile(tokenFile string) (*KiroTokenData, error) {
// Expand ~ to home directory
if strings.HasPrefix(tokenFile, "~") {
home, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("failed to get home directory: %w", err)
}
tokenFile = filepath.Join(home, tokenFile[1:])
}
data, err := os.ReadFile(tokenFile)
if err != nil {
return nil, fmt.Errorf("failed to read token file: %w", err)
}
var tokenData KiroTokenData
if err := json.Unmarshal(data, &tokenData); err != nil {
return nil, fmt.Errorf("failed to parse token file: %w", err)
}
return &tokenData, nil
}
// IsTokenExpired checks if the token has expired.
// This method parses the expiration timestamp and compares it with the current time.
//
// Parameters:
// - tokenData: The token data to check
//
// Returns:
// - bool: True if the token has expired, false otherwise
func (k *KiroAuth) IsTokenExpired(tokenData *KiroTokenData) bool {
if tokenData.ExpiresAt == "" {
return true
}
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
// Try alternate format
expiresAt, err = time.Parse("2006-01-02T15:04:05.000Z", tokenData.ExpiresAt)
if err != nil {
return true
}
}
return time.Now().After(expiresAt)
}
// makeRequest sends a request to the CodeWhisperer API.
// This is an internal method for making authenticated API calls.
//
// Parameters:
// - ctx: The context for the request
// - target: The API target (e.g., "AmazonCodeWhispererService.GetUsageLimits")
// - accessToken: The OAuth access token
// - payload: The request payload
//
// Returns:
// - []byte: The response body
// - error: An error if the request fails
func (k *KiroAuth) makeRequest(ctx context.Context, target string, accessToken string, payload interface{}) ([]byte, error) {
jsonBody, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, k.endpoint, strings.NewReader(string(jsonBody)))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
req.Header.Set("x-amz-target", target)
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := k.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
log.Errorf("failed to close response body: %v", errClose)
}
}()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body))
}
return body, nil
}
// GetUsageLimits retrieves usage information from the CodeWhisperer API.
// This method fetches the current usage statistics and subscription information.
//
// Parameters:
// - ctx: The context for the request
// - tokenData: The token data containing access token and profile ARN
//
// Returns:
// - *KiroUsageInfo: The usage information
// - error: An error if the request fails
func (k *KiroAuth) GetUsageLimits(ctx context.Context, tokenData *KiroTokenData) (*KiroUsageInfo, error) {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
"profileArn": tokenData.ProfileArn,
"resourceType": "AGENTIC_REQUEST",
}
body, err := k.makeRequest(ctx, targetGetUsage, tokenData.AccessToken, payload)
if err != nil {
return nil, err
}
var result struct {
SubscriptionInfo struct {
SubscriptionTitle string `json:"subscriptionTitle"`
} `json:"subscriptionInfo"`
UsageBreakdownList []struct {
CurrentUsageWithPrecision float64 `json:"currentUsageWithPrecision"`
UsageLimitWithPrecision float64 `json:"usageLimitWithPrecision"`
} `json:"usageBreakdownList"`
NextDateReset float64 `json:"nextDateReset"`
}
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse usage response: %w", err)
}
usage := &KiroUsageInfo{
SubscriptionTitle: result.SubscriptionInfo.SubscriptionTitle,
NextReset: fmt.Sprintf("%v", result.NextDateReset),
}
if len(result.UsageBreakdownList) > 0 {
usage.CurrentUsage = result.UsageBreakdownList[0].CurrentUsageWithPrecision
usage.UsageLimit = result.UsageBreakdownList[0].UsageLimitWithPrecision
}
return usage, nil
}
// ListAvailableModels retrieves available models from the CodeWhisperer API.
// This method fetches the list of AI models available for the authenticated user.
//
// Parameters:
// - ctx: The context for the request
// - tokenData: The token data containing access token and profile ARN
//
// Returns:
// - []*KiroModel: The list of available models
// - error: An error if the request fails
func (k *KiroAuth) ListAvailableModels(ctx context.Context, tokenData *KiroTokenData) ([]*KiroModel, error) {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
"profileArn": tokenData.ProfileArn,
}
body, err := k.makeRequest(ctx, targetListModels, tokenData.AccessToken, payload)
if err != nil {
return nil, err
}
var result struct {
Models []struct {
ModelID string `json:"modelId"`
ModelName string `json:"modelName"`
Description string `json:"description"`
RateMultiplier float64 `json:"rateMultiplier"`
RateUnit string `json:"rateUnit"`
TokenLimits struct {
MaxInputTokens int `json:"maxInputTokens"`
} `json:"tokenLimits"`
} `json:"models"`
}
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse models response: %w", err)
}
models := make([]*KiroModel, 0, len(result.Models))
for _, m := range result.Models {
models = append(models, &KiroModel{
ModelID: m.ModelID,
ModelName: m.ModelName,
Description: m.Description,
RateMultiplier: m.RateMultiplier,
RateUnit: m.RateUnit,
MaxInputTokens: m.TokenLimits.MaxInputTokens,
})
}
return models, nil
}
// CreateTokenStorage creates a new KiroTokenStorage from token data.
// This method converts the token data into a storage structure suitable for persistence.
//
// Parameters:
// - tokenData: The token data to convert
//
// Returns:
// - *KiroTokenStorage: A new token storage instance
func (k *KiroAuth) CreateTokenStorage(tokenData *KiroTokenData) *KiroTokenStorage {
return &KiroTokenStorage{
AccessToken: tokenData.AccessToken,
RefreshToken: tokenData.RefreshToken,
ProfileArn: tokenData.ProfileArn,
ExpiresAt: tokenData.ExpiresAt,
AuthMethod: tokenData.AuthMethod,
Provider: tokenData.Provider,
LastRefresh: time.Now().Format(time.RFC3339),
}
}
// ValidateToken checks if the token is valid by making a test API call.
// This method verifies the token by attempting to fetch usage limits.
//
// Parameters:
// - ctx: The context for the request
// - tokenData: The token data to validate
//
// Returns:
// - error: An error if the token is invalid
func (k *KiroAuth) ValidateToken(ctx context.Context, tokenData *KiroTokenData) error {
_, err := k.GetUsageLimits(ctx, tokenData)
return err
}
// UpdateTokenStorage updates an existing token storage with new token data.
// This method refreshes the token storage with newly obtained access and refresh tokens.
//
// Parameters:
// - storage: The existing token storage to update
// - tokenData: The new token data to apply
func (k *KiroAuth) UpdateTokenStorage(storage *KiroTokenStorage, tokenData *KiroTokenData) {
storage.AccessToken = tokenData.AccessToken
storage.RefreshToken = tokenData.RefreshToken
storage.ProfileArn = tokenData.ProfileArn
storage.ExpiresAt = tokenData.ExpiresAt
storage.AuthMethod = tokenData.AuthMethod
storage.Provider = tokenData.Provider
storage.LastRefresh = time.Now().Format(time.RFC3339)
}

View File

@@ -0,0 +1,161 @@
package kiro
import (
"encoding/base64"
"encoding/json"
"testing"
)
func TestExtractEmailFromJWT(t *testing.T) {
tests := []struct {
name string
token string
expected string
}{
{
name: "Empty token",
token: "",
expected: "",
},
{
name: "Invalid token format",
token: "not.a.valid.jwt",
expected: "",
},
{
name: "Invalid token - not base64",
token: "xxx.yyy.zzz",
expected: "",
},
{
name: "Valid JWT with email",
token: createTestJWT(map[string]any{"email": "test@example.com", "sub": "user123"}),
expected: "test@example.com",
},
{
name: "JWT without email but with preferred_username",
token: createTestJWT(map[string]any{"preferred_username": "user@domain.com", "sub": "user123"}),
expected: "user@domain.com",
},
{
name: "JWT with email-like sub",
token: createTestJWT(map[string]any{"sub": "another@test.com"}),
expected: "another@test.com",
},
{
name: "JWT without any email fields",
token: createTestJWT(map[string]any{"sub": "user123", "name": "Test User"}),
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ExtractEmailFromJWT(tt.token)
if result != tt.expected {
t.Errorf("ExtractEmailFromJWT() = %q, want %q", result, tt.expected)
}
})
}
}
func TestSanitizeEmailForFilename(t *testing.T) {
tests := []struct {
name string
email string
expected string
}{
{
name: "Empty email",
email: "",
expected: "",
},
{
name: "Simple email",
email: "user@example.com",
expected: "user@example.com",
},
{
name: "Email with space",
email: "user name@example.com",
expected: "user_name@example.com",
},
{
name: "Email with special chars",
email: "user:name@example.com",
expected: "user_name@example.com",
},
{
name: "Email with multiple special chars",
email: "user/name:test@example.com",
expected: "user_name_test@example.com",
},
{
name: "Path traversal attempt",
email: "../../../etc/passwd",
expected: "_.__.__._etc_passwd",
},
{
name: "Path traversal with backslash",
email: `..\..\..\..\windows\system32`,
expected: "_.__.__.__._windows_system32",
},
{
name: "Null byte injection attempt",
email: "user\x00@evil.com",
expected: "user_@evil.com",
},
// URL-encoded path traversal tests
{
name: "URL-encoded slash",
email: "user%2Fpath@example.com",
expected: "user_path@example.com",
},
{
name: "URL-encoded backslash",
email: "user%5Cpath@example.com",
expected: "user_path@example.com",
},
{
name: "URL-encoded dot",
email: "%2E%2E%2Fetc%2Fpasswd",
expected: "___etc_passwd",
},
{
name: "URL-encoded null",
email: "user%00@evil.com",
expected: "user_@evil.com",
},
{
name: "Double URL-encoding attack",
email: "%252F%252E%252E",
expected: "_252F_252E_252E", // % replaced with _, remaining chars preserved (safe)
},
{
name: "Mixed case URL-encoding",
email: "%2f%2F%5c%5C",
expected: "____",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := SanitizeEmailForFilename(tt.email)
if result != tt.expected {
t.Errorf("SanitizeEmailForFilename() = %q, want %q", result, tt.expected)
}
})
}
}
// createTestJWT creates a test JWT token with the given claims
func createTestJWT(claims map[string]any) string {
header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
payloadBytes, _ := json.Marshal(claims)
payload := base64.RawURLEncoding.EncodeToString(payloadBytes)
signature := base64.RawURLEncoding.EncodeToString([]byte("fake-signature"))
return header + "." + payload + "." + signature
}

296
internal/auth/kiro/oauth.go Normal file
View File

@@ -0,0 +1,296 @@
// Package kiro provides OAuth2 authentication for Kiro using native Google login.
package kiro
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"fmt"
"html"
"io"
"net"
"net/http"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
const (
// Kiro auth endpoint
kiroAuthEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
// Default callback port
defaultCallbackPort = 9876
// Auth timeout
authTimeout = 10 * time.Minute
)
// KiroTokenResponse represents the response from Kiro token endpoint.
type KiroTokenResponse struct {
AccessToken string `json:"accessToken"`
RefreshToken string `json:"refreshToken"`
ProfileArn string `json:"profileArn"`
ExpiresIn int `json:"expiresIn"`
}
// KiroOAuth handles the OAuth flow for Kiro authentication.
type KiroOAuth struct {
httpClient *http.Client
cfg *config.Config
}
// NewKiroOAuth creates a new Kiro OAuth handler.
func NewKiroOAuth(cfg *config.Config) *KiroOAuth {
client := &http.Client{Timeout: 30 * time.Second}
if cfg != nil {
client = util.SetProxy(&cfg.SDKConfig, client)
}
return &KiroOAuth{
httpClient: client,
cfg: cfg,
}
}
// generateCodeVerifier generates a random code verifier for PKCE.
func generateCodeVerifier() (string, error) {
b := make([]byte, 32)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// generateCodeChallenge generates the code challenge from verifier.
func generateCodeChallenge(verifier string) string {
h := sha256.Sum256([]byte(verifier))
return base64.RawURLEncoding.EncodeToString(h[:])
}
// generateState generates a random state parameter.
func generateState() (string, error) {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// AuthResult contains the authorization code and state from callback.
type AuthResult struct {
Code string
State string
Error string
}
// startCallbackServer starts a local HTTP server to receive the OAuth callback.
func (o *KiroOAuth) startCallbackServer(ctx context.Context, expectedState string) (string, <-chan AuthResult, error) {
// Try to find an available port - use localhost like Kiro does
listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", defaultCallbackPort))
if err != nil {
// Try with dynamic port (RFC 8252 allows dynamic ports for native apps)
log.Warnf("kiro oauth: default port %d is busy, falling back to dynamic port", defaultCallbackPort)
listener, err = net.Listen("tcp", "localhost:0")
if err != nil {
return "", nil, fmt.Errorf("failed to start callback server: %w", err)
}
}
port := listener.Addr().(*net.TCPAddr).Port
// Use http scheme for local callback server
redirectURI := fmt.Sprintf("http://localhost:%d/oauth/callback", port)
resultChan := make(chan AuthResult, 1)
server := &http.Server{
ReadHeaderTimeout: 10 * time.Second,
}
mux := http.NewServeMux()
mux.HandleFunc("/oauth/callback", func(w http.ResponseWriter, r *http.Request) {
code := r.URL.Query().Get("code")
state := r.URL.Query().Get("state")
errParam := r.URL.Query().Get("error")
if errParam != "" {
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, `<html><body><h1>Login Failed</h1><p>%s</p><p>You can close this window.</p></body></html>`, html.EscapeString(errParam))
resultChan <- AuthResult{Error: errParam}
return
}
if state != expectedState {
w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusBadRequest)
fmt.Fprint(w, `<html><body><h1>Login Failed</h1><p>Invalid state parameter</p><p>You can close this window.</p></body></html>`)
resultChan <- AuthResult{Error: "state mismatch"}
return
}
w.Header().Set("Content-Type", "text/html")
fmt.Fprint(w, `<html><body><h1>Login Successful!</h1><p>You can close this window and return to the terminal.</p></body></html>`)
resultChan <- AuthResult{Code: code, State: state}
})
server.Handler = mux
go func() {
if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
log.Debugf("callback server error: %v", err)
}
}()
go func() {
select {
case <-ctx.Done():
case <-time.After(authTimeout):
case <-resultChan:
}
_ = server.Shutdown(context.Background())
}()
return redirectURI, resultChan, nil
}
// LoginWithBuilderID performs OAuth login with AWS Builder ID using device code flow.
func (o *KiroOAuth) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
ssoClient := NewSSOOIDCClient(o.cfg)
return ssoClient.LoginWithBuilderID(ctx)
}
// exchangeCodeForToken exchanges the authorization code for tokens.
func (o *KiroOAuth) exchangeCodeForToken(ctx context.Context, code, codeVerifier, redirectURI string) (*KiroTokenData, error) {
payload := map[string]string{
"code": code,
"code_verifier": codeVerifier,
"redirect_uri": redirectURI,
}
body, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
tokenURL := kiroAuthEndpoint + "/oauth/token"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := o.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("token request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
}
var tokenResp KiroTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse token response: %w", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: "", // Caller should preserve original provider
}, nil
}
// RefreshToken refreshes an expired access token.
func (o *KiroOAuth) RefreshToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
payload := map[string]string{
"refreshToken": refreshToken,
}
body, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
refreshURL := kiroAuthEndpoint + "/refreshToken"
req, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := o.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("refresh request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
}
var tokenResp KiroTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse token response: %w", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: "", // Caller should preserve original provider
}, nil
}
// LoginWithGoogle performs OAuth login with Google using Kiro's social auth.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (o *KiroOAuth) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
socialClient := NewSocialAuthClient(o.cfg)
return socialClient.LoginWithGoogle(ctx)
}
// LoginWithGitHub performs OAuth login with GitHub using Kiro's social auth.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (o *KiroOAuth) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
socialClient := NewSocialAuthClient(o.cfg)
return socialClient.LoginWithGitHub(ctx)
}

View File

@@ -0,0 +1,725 @@
// Package kiro provides custom protocol handler registration for Kiro OAuth.
// This enables the CLI to intercept kiro:// URIs for social authentication (Google/GitHub).
package kiro
import (
"context"
"fmt"
"html"
"net"
"net/http"
"net/url"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
)
const (
// KiroProtocol is the custom URI scheme used by Kiro
KiroProtocol = "kiro"
// KiroAuthority is the URI authority for authentication callbacks
KiroAuthority = "kiro.kiroAgent"
// KiroAuthPath is the path for successful authentication
KiroAuthPath = "/authenticate-success"
// KiroRedirectURI is the full redirect URI for social auth
KiroRedirectURI = "kiro://kiro.kiroAgent/authenticate-success"
// DefaultHandlerPort is the default port for the local callback server
DefaultHandlerPort = 19876
// HandlerTimeout is how long to wait for the OAuth callback
HandlerTimeout = 10 * time.Minute
)
// ProtocolHandler manages the custom kiro:// protocol handler for OAuth callbacks.
type ProtocolHandler struct {
port int
server *http.Server
listener net.Listener
resultChan chan *AuthCallback
stopChan chan struct{}
mu sync.Mutex
running bool
}
// AuthCallback contains the OAuth callback parameters.
type AuthCallback struct {
Code string
State string
Error string
}
// NewProtocolHandler creates a new protocol handler.
func NewProtocolHandler() *ProtocolHandler {
return &ProtocolHandler{
port: DefaultHandlerPort,
resultChan: make(chan *AuthCallback, 1),
stopChan: make(chan struct{}),
}
}
// Start starts the local callback server that receives redirects from the protocol handler.
func (h *ProtocolHandler) Start(ctx context.Context) (int, error) {
h.mu.Lock()
defer h.mu.Unlock()
if h.running {
return h.port, nil
}
// Drain any stale results from previous runs
select {
case <-h.resultChan:
default:
}
// Reset stopChan for reuse - close old channel first to unblock any waiting goroutines
if h.stopChan != nil {
select {
case <-h.stopChan:
// Already closed
default:
close(h.stopChan)
}
}
h.stopChan = make(chan struct{})
// Try ports in known range (must match handler script port range)
var listener net.Listener
var err error
portRange := []int{DefaultHandlerPort, DefaultHandlerPort + 1, DefaultHandlerPort + 2, DefaultHandlerPort + 3, DefaultHandlerPort + 4}
for _, port := range portRange {
listener, err = net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
if err == nil {
break
}
log.Debugf("kiro protocol handler: port %d busy, trying next", port)
}
if listener == nil {
return 0, fmt.Errorf("failed to start callback server: all ports %d-%d are busy", DefaultHandlerPort, DefaultHandlerPort+4)
}
h.listener = listener
h.port = listener.Addr().(*net.TCPAddr).Port
mux := http.NewServeMux()
mux.HandleFunc("/oauth/callback", h.handleCallback)
h.server = &http.Server{
Handler: mux,
ReadHeaderTimeout: 10 * time.Second,
}
go func() {
if err := h.server.Serve(listener); err != nil && err != http.ErrServerClosed {
log.Debugf("kiro protocol handler server error: %v", err)
}
}()
h.running = true
log.Debugf("kiro protocol handler started on port %d", h.port)
// Auto-shutdown after context done, timeout, or explicit stop
// Capture references to prevent race with new Start() calls
currentStopChan := h.stopChan
currentServer := h.server
currentListener := h.listener
go func() {
select {
case <-ctx.Done():
case <-time.After(HandlerTimeout):
case <-currentStopChan:
return // Already stopped, exit goroutine
}
// Only stop if this is still the current server/listener instance
h.mu.Lock()
if h.server == currentServer && h.listener == currentListener {
h.mu.Unlock()
h.Stop()
} else {
h.mu.Unlock()
}
}()
return h.port, nil
}
// Stop stops the callback server.
func (h *ProtocolHandler) Stop() {
h.mu.Lock()
defer h.mu.Unlock()
if !h.running {
return
}
// Signal the auto-shutdown goroutine to exit.
// This select pattern is safe because stopChan is only modified while holding h.mu,
// and we hold the lock here. The select prevents panic from double-close.
select {
case <-h.stopChan:
// Already closed
default:
close(h.stopChan)
}
if h.server != nil {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_ = h.server.Shutdown(ctx)
}
h.running = false
log.Debug("kiro protocol handler stopped")
}
// WaitForCallback waits for the OAuth callback and returns the result.
func (h *ProtocolHandler) WaitForCallback(ctx context.Context) (*AuthCallback, error) {
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(HandlerTimeout):
return nil, fmt.Errorf("timeout waiting for OAuth callback")
case result := <-h.resultChan:
return result, nil
}
}
// GetPort returns the port the handler is listening on.
func (h *ProtocolHandler) GetPort() int {
return h.port
}
// handleCallback processes the OAuth callback from the protocol handler script.
func (h *ProtocolHandler) handleCallback(w http.ResponseWriter, r *http.Request) {
code := r.URL.Query().Get("code")
state := r.URL.Query().Get("state")
errParam := r.URL.Query().Get("error")
result := &AuthCallback{
Code: code,
State: state,
Error: errParam,
}
// Send result
select {
case h.resultChan <- result:
default:
// Channel full, ignore duplicate callbacks
}
// Send success response
w.Header().Set("Content-Type", "text/html; charset=utf-8")
if errParam != "" {
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, `<!DOCTYPE html>
<html>
<head><title>Login Failed</title></head>
<body>
<h1>Login Failed</h1>
<p>Error: %s</p>
<p>You can close this window.</p>
</body>
</html>`, html.EscapeString(errParam))
} else {
fmt.Fprint(w, `<!DOCTYPE html>
<html>
<head><title>Login Successful</title></head>
<body>
<h1>Login Successful!</h1>
<p>You can close this window and return to the terminal.</p>
<script>window.close();</script>
</body>
</html>`)
}
}
// IsProtocolHandlerInstalled checks if the kiro:// protocol handler is installed.
func IsProtocolHandlerInstalled() bool {
switch runtime.GOOS {
case "linux":
return isLinuxHandlerInstalled()
case "windows":
return isWindowsHandlerInstalled()
case "darwin":
return isDarwinHandlerInstalled()
default:
return false
}
}
// InstallProtocolHandler installs the kiro:// protocol handler for the current platform.
func InstallProtocolHandler(handlerPort int) error {
switch runtime.GOOS {
case "linux":
return installLinuxHandler(handlerPort)
case "windows":
return installWindowsHandler(handlerPort)
case "darwin":
return installDarwinHandler(handlerPort)
default:
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
}
}
// UninstallProtocolHandler removes the kiro:// protocol handler.
func UninstallProtocolHandler() error {
switch runtime.GOOS {
case "linux":
return uninstallLinuxHandler()
case "windows":
return uninstallWindowsHandler()
case "darwin":
return uninstallDarwinHandler()
default:
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
}
}
// --- Linux Implementation ---
func getLinuxDesktopPath() string {
homeDir, _ := os.UserHomeDir()
return filepath.Join(homeDir, ".local", "share", "applications", "kiro-oauth-handler.desktop")
}
func getLinuxHandlerScriptPath() string {
homeDir, _ := os.UserHomeDir()
return filepath.Join(homeDir, ".local", "bin", "kiro-oauth-handler")
}
func isLinuxHandlerInstalled() bool {
desktopPath := getLinuxDesktopPath()
_, err := os.Stat(desktopPath)
return err == nil
}
func installLinuxHandler(handlerPort int) error {
// Create directories
homeDir, err := os.UserHomeDir()
if err != nil {
return err
}
binDir := filepath.Join(homeDir, ".local", "bin")
appDir := filepath.Join(homeDir, ".local", "share", "applications")
if err := os.MkdirAll(binDir, 0755); err != nil {
return fmt.Errorf("failed to create bin directory: %w", err)
}
if err := os.MkdirAll(appDir, 0755); err != nil {
return fmt.Errorf("failed to create applications directory: %w", err)
}
// Create handler script - tries multiple ports to handle dynamic port allocation
scriptPath := getLinuxHandlerScriptPath()
scriptContent := fmt.Sprintf(`#!/bin/bash
# Kiro OAuth Protocol Handler
# Handles kiro:// URIs - tries CLI first, then forwards to Kiro IDE
URL="$1"
# Check curl availability
if ! command -v curl &> /dev/null; then
echo "Error: curl is required for Kiro OAuth handler" >&2
exit 1
fi
# Extract code and state from URL
[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
# Try CLI proxy on multiple possible ports (default + dynamic range)
CLI_OK=0
for PORT in %d %d %d %d %d; do
if [ -n "$ERROR" ]; then
curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && CLI_OK=1 && break
elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && CLI_OK=1 && break
fi
done
# If CLI not available, forward to Kiro IDE
if [ $CLI_OK -eq 0 ] && [ -x "/usr/share/kiro/kiro" ]; then
/usr/share/kiro/kiro --open-url "$URL" &
fi
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
if err := os.WriteFile(scriptPath, []byte(scriptContent), 0755); err != nil {
return fmt.Errorf("failed to write handler script: %w", err)
}
// Create .desktop file
desktopPath := getLinuxDesktopPath()
desktopContent := fmt.Sprintf(`[Desktop Entry]
Name=Kiro OAuth Handler
Comment=Handle kiro:// protocol for CLI Proxy API authentication
Exec=%s %%u
Type=Application
Terminal=false
NoDisplay=true
MimeType=x-scheme-handler/kiro;
Categories=Utility;
`, scriptPath)
if err := os.WriteFile(desktopPath, []byte(desktopContent), 0644); err != nil {
return fmt.Errorf("failed to write desktop file: %w", err)
}
// Register handler with xdg-mime
cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
if err := cmd.Run(); err != nil {
log.Warnf("xdg-mime registration failed (may need manual setup): %v", err)
}
// Update desktop database
cmd = exec.Command("update-desktop-database", appDir)
_ = cmd.Run() // Ignore errors, not critical
log.Info("Kiro protocol handler installed for Linux")
return nil
}
func uninstallLinuxHandler() error {
desktopPath := getLinuxDesktopPath()
scriptPath := getLinuxHandlerScriptPath()
if err := os.Remove(desktopPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove desktop file: %w", err)
}
if err := os.Remove(scriptPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove handler script: %w", err)
}
log.Info("Kiro protocol handler uninstalled")
return nil
}
// --- Windows Implementation ---
func isWindowsHandlerInstalled() bool {
// Check registry key existence
cmd := exec.Command("reg", "query", `HKCU\Software\Classes\kiro`, "/ve")
return cmd.Run() == nil
}
func installWindowsHandler(handlerPort int) error {
homeDir, err := os.UserHomeDir()
if err != nil {
return err
}
// Create handler script (PowerShell)
scriptDir := filepath.Join(homeDir, ".cliproxyapi")
if err := os.MkdirAll(scriptDir, 0755); err != nil {
return fmt.Errorf("failed to create script directory: %w", err)
}
scriptPath := filepath.Join(scriptDir, "kiro-oauth-handler.ps1")
scriptContent := fmt.Sprintf(`# Kiro OAuth Protocol Handler for Windows
param([string]$url)
# Load required assembly for HttpUtility
Add-Type -AssemblyName System.Web
# Parse URL parameters
$uri = [System.Uri]$url
$query = [System.Web.HttpUtility]::ParseQueryString($uri.Query)
$code = $query["code"]
$state = $query["state"]
$errorParam = $query["error"]
# Try multiple ports (default + dynamic range)
$ports = @(%d, %d, %d, %d, %d)
$success = $false
foreach ($port in $ports) {
if ($success) { break }
$callbackUrl = "http://127.0.0.1:$port/oauth/callback"
try {
if ($errorParam) {
$fullUrl = $callbackUrl + "?error=" + $errorParam
Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
$success = $true
} elseif ($code -and $state) {
$fullUrl = $callbackUrl + "?code=" + $code + "&state=" + $state
Invoke-WebRequest -Uri $fullUrl -UseBasicParsing -TimeoutSec 1 -ErrorAction Stop | Out-Null
$success = $true
}
} catch {
# Try next port
}
}
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
if err := os.WriteFile(scriptPath, []byte(scriptContent), 0644); err != nil {
return fmt.Errorf("failed to write handler script: %w", err)
}
// Create batch wrapper
batchPath := filepath.Join(scriptDir, "kiro-oauth-handler.bat")
batchContent := fmt.Sprintf("@echo off\npowershell -ExecutionPolicy Bypass -File \"%s\" \"%%1\"\n", scriptPath)
if err := os.WriteFile(batchPath, []byte(batchContent), 0644); err != nil {
return fmt.Errorf("failed to write batch wrapper: %w", err)
}
// Register in Windows registry
commands := [][]string{
{"reg", "add", `HKCU\Software\Classes\kiro`, "/ve", "/d", "URL:Kiro Protocol", "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro`, "/v", "URL Protocol", "/d", "", "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro\shell`, "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro\shell\open`, "/f"},
{"reg", "add", `HKCU\Software\Classes\kiro\shell\open\command`, "/ve", "/d", fmt.Sprintf("\"%s\" \"%%1\"", batchPath), "/f"},
}
for _, args := range commands {
cmd := exec.Command(args[0], args[1:]...)
if err := cmd.Run(); err != nil {
return fmt.Errorf("failed to run registry command: %w", err)
}
}
log.Info("Kiro protocol handler installed for Windows")
return nil
}
func uninstallWindowsHandler() error {
// Remove registry keys
cmd := exec.Command("reg", "delete", `HKCU\Software\Classes\kiro`, "/f")
if err := cmd.Run(); err != nil {
log.Warnf("failed to remove registry key: %v", err)
}
// Remove scripts
homeDir, _ := os.UserHomeDir()
scriptDir := filepath.Join(homeDir, ".cliproxyapi")
_ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.ps1"))
_ = os.Remove(filepath.Join(scriptDir, "kiro-oauth-handler.bat"))
log.Info("Kiro protocol handler uninstalled")
return nil
}
// --- macOS Implementation ---
func getDarwinAppPath() string {
homeDir, _ := os.UserHomeDir()
return filepath.Join(homeDir, "Applications", "KiroOAuthHandler.app")
}
func isDarwinHandlerInstalled() bool {
appPath := getDarwinAppPath()
_, err := os.Stat(appPath)
return err == nil
}
func installDarwinHandler(handlerPort int) error {
// Create app bundle structure
appPath := getDarwinAppPath()
contentsPath := filepath.Join(appPath, "Contents")
macOSPath := filepath.Join(contentsPath, "MacOS")
if err := os.MkdirAll(macOSPath, 0755); err != nil {
return fmt.Errorf("failed to create app bundle: %w", err)
}
// Create Info.plist
plistPath := filepath.Join(contentsPath, "Info.plist")
plistContent := `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>CFBundleIdentifier</key>
<string>com.cliproxyapi.kiro-oauth-handler</string>
<key>CFBundleName</key>
<string>KiroOAuthHandler</string>
<key>CFBundleExecutable</key>
<string>kiro-oauth-handler</string>
<key>CFBundleVersion</key>
<string>1.0</string>
<key>CFBundleURLTypes</key>
<array>
<dict>
<key>CFBundleURLName</key>
<string>Kiro Protocol</string>
<key>CFBundleURLSchemes</key>
<array>
<string>kiro</string>
</array>
</dict>
</array>
<key>LSBackgroundOnly</key>
<true/>
</dict>
</plist>`
if err := os.WriteFile(plistPath, []byte(plistContent), 0644); err != nil {
return fmt.Errorf("failed to write Info.plist: %w", err)
}
// Create executable script - tries multiple ports to handle dynamic port allocation
execPath := filepath.Join(macOSPath, "kiro-oauth-handler")
execContent := fmt.Sprintf(`#!/bin/bash
# Kiro OAuth Protocol Handler for macOS
URL="$1"
# Check curl availability (should always exist on macOS)
if [ ! -x /usr/bin/curl ]; then
echo "Error: curl is required for Kiro OAuth handler" >&2
exit 1
fi
# Extract code and state from URL
[[ "$URL" =~ code=([^&]+) ]] && CODE="${BASH_REMATCH[1]}"
[[ "$URL" =~ state=([^&]+) ]] && STATE="${BASH_REMATCH[1]}"
[[ "$URL" =~ error=([^&]+) ]] && ERROR="${BASH_REMATCH[1]}"
# Try multiple ports (default + dynamic range)
for PORT in %d %d %d %d %d; do
if [ -n "$ERROR" ]; then
/usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?error=$ERROR" && exit 0
elif [ -n "$CODE" ] && [ -n "$STATE" ]; then
/usr/bin/curl -sf --connect-timeout 1 "http://127.0.0.1:$PORT/oauth/callback?code=$CODE&state=$STATE" && exit 0
fi
done
`, handlerPort, handlerPort+1, handlerPort+2, handlerPort+3, handlerPort+4)
if err := os.WriteFile(execPath, []byte(execContent), 0755); err != nil {
return fmt.Errorf("failed to write executable: %w", err)
}
// Register the app with Launch Services
cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
"-f", appPath)
if err := cmd.Run(); err != nil {
log.Warnf("lsregister failed (handler may still work): %v", err)
}
log.Info("Kiro protocol handler installed for macOS")
return nil
}
func uninstallDarwinHandler() error {
appPath := getDarwinAppPath()
// Unregister from Launch Services
cmd := exec.Command("/System/Library/Frameworks/CoreServices.framework/Frameworks/LaunchServices.framework/Support/lsregister",
"-u", appPath)
_ = cmd.Run()
// Remove app bundle
if err := os.RemoveAll(appPath); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to remove app bundle: %w", err)
}
log.Info("Kiro protocol handler uninstalled")
return nil
}
// ParseKiroURI parses a kiro:// URI and extracts the callback parameters.
func ParseKiroURI(rawURI string) (*AuthCallback, error) {
u, err := url.Parse(rawURI)
if err != nil {
return nil, fmt.Errorf("invalid URI: %w", err)
}
if u.Scheme != KiroProtocol {
return nil, fmt.Errorf("invalid scheme: expected %s, got %s", KiroProtocol, u.Scheme)
}
if u.Host != KiroAuthority {
return nil, fmt.Errorf("invalid authority: expected %s, got %s", KiroAuthority, u.Host)
}
query := u.Query()
return &AuthCallback{
Code: query.Get("code"),
State: query.Get("state"),
Error: query.Get("error"),
}, nil
}
// GetHandlerInstructions returns platform-specific instructions for manual handler setup.
func GetHandlerInstructions() string {
switch runtime.GOOS {
case "linux":
return `To manually set up the Kiro protocol handler on Linux:
1. Create ~/.local/share/applications/kiro-oauth-handler.desktop:
[Desktop Entry]
Name=Kiro OAuth Handler
Exec=~/.local/bin/kiro-oauth-handler %u
Type=Application
Terminal=false
MimeType=x-scheme-handler/kiro;
2. Create ~/.local/bin/kiro-oauth-handler (make it executable):
#!/bin/bash
URL="$1"
# ... (see generated script for full content)
3. Run: xdg-mime default kiro-oauth-handler.desktop x-scheme-handler/kiro`
case "windows":
return `To manually set up the Kiro protocol handler on Windows:
1. Open Registry Editor (regedit.exe)
2. Create key: HKEY_CURRENT_USER\Software\Classes\kiro
3. Set default value to: URL:Kiro Protocol
4. Create string value "URL Protocol" with empty data
5. Create subkey: shell\open\command
6. Set default value to: "C:\path\to\handler.bat" "%1"`
case "darwin":
return `To manually set up the Kiro protocol handler on macOS:
1. Create ~/Applications/KiroOAuthHandler.app bundle
2. Add Info.plist with CFBundleURLTypes containing "kiro" scheme
3. Create executable in Contents/MacOS/
4. Run: /System/Library/.../lsregister -f ~/Applications/KiroOAuthHandler.app`
default:
return "Protocol handler setup is not supported on this platform."
}
}
// SetupProtocolHandlerIfNeeded checks and installs the protocol handler if needed.
func SetupProtocolHandlerIfNeeded(handlerPort int) error {
if IsProtocolHandlerInstalled() {
log.Debug("Kiro protocol handler already installed")
return nil
}
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
fmt.Println("║ Kiro Protocol Handler Setup Required ║")
fmt.Println("╚══════════════════════════════════════════════════════════╝")
fmt.Println("\nTo enable Google/GitHub login, we need to install a protocol handler.")
fmt.Println("This allows your browser to redirect back to the CLI after authentication.")
fmt.Println("\nInstalling protocol handler...")
if err := InstallProtocolHandler(handlerPort); err != nil {
fmt.Printf("\n⚠ Automatic installation failed: %v\n", err)
fmt.Println("\nManual setup instructions:")
fmt.Println(strings.Repeat("-", 60))
fmt.Println(GetHandlerInstructions())
return err
}
fmt.Println("\n✓ Protocol handler installed successfully!")
return nil
}

View File

@@ -0,0 +1,403 @@
// Package kiro provides social authentication (Google/GitHub) for Kiro via AuthServiceClient.
package kiro
import (
"bufio"
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"os"
"os/exec"
"runtime"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
"golang.org/x/term"
)
const (
// Kiro AuthService endpoint
kiroAuthServiceEndpoint = "https://prod.us-east-1.auth.desktop.kiro.dev"
// OAuth timeout
socialAuthTimeout = 10 * time.Minute
)
// SocialProvider represents the social login provider.
type SocialProvider string
const (
// ProviderGoogle is Google OAuth provider
ProviderGoogle SocialProvider = "Google"
// ProviderGitHub is GitHub OAuth provider
ProviderGitHub SocialProvider = "Github"
// Note: AWS Builder ID is NOT supported by Kiro's auth service.
// It only supports: Google, Github, Cognito
// AWS Builder ID must use device code flow via SSO OIDC.
)
// CreateTokenRequest is sent to Kiro's /oauth/token endpoint.
type CreateTokenRequest struct {
Code string `json:"code"`
CodeVerifier string `json:"code_verifier"`
RedirectURI string `json:"redirect_uri"`
InvitationCode string `json:"invitation_code,omitempty"`
}
// SocialTokenResponse from Kiro's /oauth/token endpoint for social auth.
type SocialTokenResponse struct {
AccessToken string `json:"accessToken"`
RefreshToken string `json:"refreshToken"`
ProfileArn string `json:"profileArn"`
ExpiresIn int `json:"expiresIn"`
}
// RefreshTokenRequest is sent to Kiro's /refreshToken endpoint.
type RefreshTokenRequest struct {
RefreshToken string `json:"refreshToken"`
}
// SocialAuthClient handles social authentication with Kiro.
type SocialAuthClient struct {
httpClient *http.Client
cfg *config.Config
protocolHandler *ProtocolHandler
}
// NewSocialAuthClient creates a new social auth client.
func NewSocialAuthClient(cfg *config.Config) *SocialAuthClient {
client := &http.Client{Timeout: 30 * time.Second}
if cfg != nil {
client = util.SetProxy(&cfg.SDKConfig, client)
}
return &SocialAuthClient{
httpClient: client,
cfg: cfg,
protocolHandler: NewProtocolHandler(),
}
}
// generatePKCE generates PKCE code verifier and challenge.
func generatePKCE() (verifier, challenge string, err error) {
// Generate 32 bytes of random data for verifier
b := make([]byte, 32)
if _, err := rand.Read(b); err != nil {
return "", "", fmt.Errorf("failed to generate random bytes: %w", err)
}
verifier = base64.RawURLEncoding.EncodeToString(b)
// Generate SHA256 hash of verifier for challenge
h := sha256.Sum256([]byte(verifier))
challenge = base64.RawURLEncoding.EncodeToString(h[:])
return verifier, challenge, nil
}
// generateState generates a random state parameter.
func generateStateParam() (string, error) {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// buildLoginURL constructs the Kiro OAuth login URL.
// The login endpoint expects a GET request with query parameters.
// Format: /login?idp=Google&redirect_uri=...&code_challenge=...&code_challenge_method=S256&state=...&prompt=select_account
// The prompt=select_account parameter forces the account selection screen even if already logged in.
func (c *SocialAuthClient) buildLoginURL(provider, redirectURI, codeChallenge, state string) string {
return fmt.Sprintf("%s/login?idp=%s&redirect_uri=%s&code_challenge=%s&code_challenge_method=S256&state=%s&prompt=select_account",
kiroAuthServiceEndpoint,
provider,
url.QueryEscape(redirectURI),
codeChallenge,
state,
)
}
// createToken exchanges the authorization code for tokens.
func (c *SocialAuthClient) createToken(ctx context.Context, req *CreateTokenRequest) (*SocialTokenResponse, error) {
body, err := json.Marshal(req)
if err != nil {
return nil, fmt.Errorf("failed to marshal token request: %w", err)
}
tokenURL := kiroAuthServiceEndpoint + "/oauth/token"
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create token request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := c.httpClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("token request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read token response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token exchange failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token exchange failed (status %d)", resp.StatusCode)
}
var tokenResp SocialTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse token response: %w", err)
}
return &tokenResp, nil
}
// RefreshSocialToken refreshes an expired social auth token.
func (c *SocialAuthClient) RefreshSocialToken(ctx context.Context, refreshToken string) (*KiroTokenData, error) {
body, err := json.Marshal(&RefreshTokenRequest{RefreshToken: refreshToken})
if err != nil {
return nil, fmt.Errorf("failed to marshal refresh request: %w", err)
}
refreshURL := kiroAuthServiceEndpoint + "/refreshToken"
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, strings.NewReader(string(body)))
if err != nil {
return nil, fmt.Errorf("failed to create refresh request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("User-Agent", "cli-proxy-api/1.0.0")
resp, err := c.httpClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("refresh request failed: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read refresh response: %w", err)
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
}
var tokenResp SocialTokenResponse
if err := json.Unmarshal(respBody, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to parse refresh response: %w", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600 // Default 1 hour
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: "", // Caller should preserve original provider
}, nil
}
// LoginWithSocial performs OAuth login with Google.
func (c *SocialAuthClient) LoginWithSocial(ctx context.Context, provider SocialProvider) (*KiroTokenData, error) {
providerName := string(provider)
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
fmt.Printf("║ Kiro Authentication (%s) ║\n", providerName)
fmt.Println("╚══════════════════════════════════════════════════════════╝")
// Step 1: Setup protocol handler
fmt.Println("\nSetting up authentication...")
// Start the local callback server
handlerPort, err := c.protocolHandler.Start(ctx)
if err != nil {
return nil, fmt.Errorf("failed to start callback server: %w", err)
}
defer c.protocolHandler.Stop()
// Ensure protocol handler is installed and set as default
if err := SetupProtocolHandlerIfNeeded(handlerPort); err != nil {
fmt.Println("\n⚠ Protocol handler setup failed. Trying alternative method...")
fmt.Println(" If you see a browser 'Open with' dialog, select your default browser.")
fmt.Println(" For manual setup instructions, run: cliproxy kiro --help-protocol")
log.Debugf("kiro: protocol handler setup error: %v", err)
// Continue anyway - user might have set it up manually or select browser manually
} else {
// Force set our handler as default (prevents "Open with" dialog)
forceDefaultProtocolHandler()
}
// Step 2: Generate PKCE codes
codeVerifier, codeChallenge, err := generatePKCE()
if err != nil {
return nil, fmt.Errorf("failed to generate PKCE: %w", err)
}
// Step 3: Generate state
state, err := generateStateParam()
if err != nil {
return nil, fmt.Errorf("failed to generate state: %w", err)
}
// Step 4: Build the login URL (Kiro uses GET request with query params)
authURL := c.buildLoginURL(providerName, KiroRedirectURI, codeChallenge, state)
// Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
// Incognito mode enables multi-account support by bypassing cached sessions
if c.cfg != nil {
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
if !c.cfg.IncognitoBrowser {
log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
} else {
log.Debug("kiro: using incognito mode for multi-account support")
}
} else {
browser.SetIncognitoMode(true) // Default to incognito if no config
log.Debug("kiro: using incognito mode for multi-account support (default)")
}
// Step 5: Open browser for user authentication
fmt.Println("\n════════════════════════════════════════════════════════════")
fmt.Printf(" Opening browser for %s authentication...\n", providerName)
fmt.Println("════════════════════════════════════════════════════════════")
fmt.Printf("\n URL: %s\n\n", authURL)
if err := browser.OpenURL(authURL); err != nil {
log.Warnf("Could not open browser automatically: %v", err)
fmt.Println(" ⚠ Could not open browser automatically.")
fmt.Println(" Please open the URL above in your browser manually.")
} else {
fmt.Println(" (Browser opened automatically)")
}
fmt.Println("\n Waiting for authentication callback...")
// Step 6: Wait for callback
callback, err := c.protocolHandler.WaitForCallback(ctx)
if err != nil {
return nil, fmt.Errorf("failed to receive callback: %w", err)
}
if callback.Error != "" {
return nil, fmt.Errorf("authentication error: %s", callback.Error)
}
if callback.State != state {
// Log state values for debugging, but don't expose in user-facing error
log.Debugf("kiro: OAuth state mismatch - expected %s, got %s", state, callback.State)
return nil, fmt.Errorf("OAuth state validation failed - please try again")
}
if callback.Code == "" {
return nil, fmt.Errorf("no authorization code received")
}
fmt.Println("\n✓ Authorization received!")
// Step 7: Exchange code for tokens
fmt.Println("Exchanging code for tokens...")
tokenReq := &CreateTokenRequest{
Code: callback.Code,
CodeVerifier: codeVerifier,
RedirectURI: KiroRedirectURI,
}
tokenResp, err := c.createToken(ctx, tokenReq)
if err != nil {
return nil, fmt.Errorf("failed to exchange code for tokens: %w", err)
}
fmt.Println("\n✓ Authentication successful!")
// Close the browser window
if err := browser.CloseBrowser(); err != nil {
log.Debugf("Failed to close browser: %v", err)
}
// Validate ExpiresIn - use default 1 hour if invalid
expiresIn := tokenResp.ExpiresIn
if expiresIn <= 0 {
expiresIn = 3600
}
expiresAt := time.Now().Add(time.Duration(expiresIn) * time.Second)
// Try to extract email from JWT access token first
email := ExtractEmailFromJWT(tokenResp.AccessToken)
// If no email in JWT, ask user for account label (only in interactive mode)
if email == "" && isInteractiveTerminal() {
fmt.Print("\n Enter account label for file naming (optional, press Enter to skip): ")
reader := bufio.NewReader(os.Stdin)
var err error
email, err = reader.ReadString('\n')
if err != nil {
log.Debugf("Failed to read account label: %v", err)
}
email = strings.TrimSpace(email)
}
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: tokenResp.ProfileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "social",
Provider: providerName,
Email: email, // JWT email or user-provided label
}, nil
}
// LoginWithGoogle performs OAuth login with Google.
func (c *SocialAuthClient) LoginWithGoogle(ctx context.Context) (*KiroTokenData, error) {
return c.LoginWithSocial(ctx, ProviderGoogle)
}
// LoginWithGitHub performs OAuth login with GitHub.
func (c *SocialAuthClient) LoginWithGitHub(ctx context.Context) (*KiroTokenData, error) {
return c.LoginWithSocial(ctx, ProviderGitHub)
}
// forceDefaultProtocolHandler sets our protocol handler as the default for kiro:// URLs.
// This prevents the "Open with" dialog from appearing on Linux.
// On non-Linux platforms, this is a no-op as they use different mechanisms.
func forceDefaultProtocolHandler() {
if runtime.GOOS != "linux" {
return // Non-Linux platforms use different handler mechanisms
}
// Set our handler as default using xdg-mime
cmd := exec.Command("xdg-mime", "default", "kiro-oauth-handler.desktop", "x-scheme-handler/kiro")
if err := cmd.Run(); err != nil {
log.Warnf("Failed to set default protocol handler: %v. You may see a handler selection dialog.", err)
}
}
// isInteractiveTerminal checks if stdin is connected to an interactive terminal.
// Returns false in CI/automated environments or when stdin is piped.
func isInteractiveTerminal() bool {
return term.IsTerminal(int(os.Stdin.Fd()))
}

View File

@@ -0,0 +1,527 @@
// Package kiro provides AWS SSO OIDC authentication for Kiro.
package kiro
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/browser"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus"
)
const (
// AWS SSO OIDC endpoints
ssoOIDCEndpoint = "https://oidc.us-east-1.amazonaws.com"
// Kiro's start URL for Builder ID
builderIDStartURL = "https://view.awsapps.com/start"
// Polling interval
pollInterval = 5 * time.Second
)
// SSOOIDCClient handles AWS SSO OIDC authentication.
type SSOOIDCClient struct {
httpClient *http.Client
cfg *config.Config
}
// NewSSOOIDCClient creates a new SSO OIDC client.
func NewSSOOIDCClient(cfg *config.Config) *SSOOIDCClient {
client := &http.Client{Timeout: 30 * time.Second}
if cfg != nil {
client = util.SetProxy(&cfg.SDKConfig, client)
}
return &SSOOIDCClient{
httpClient: client,
cfg: cfg,
}
}
// RegisterClientResponse from AWS SSO OIDC.
type RegisterClientResponse struct {
ClientID string `json:"clientId"`
ClientSecret string `json:"clientSecret"`
ClientIDIssuedAt int64 `json:"clientIdIssuedAt"`
ClientSecretExpiresAt int64 `json:"clientSecretExpiresAt"`
}
// StartDeviceAuthResponse from AWS SSO OIDC.
type StartDeviceAuthResponse struct {
DeviceCode string `json:"deviceCode"`
UserCode string `json:"userCode"`
VerificationURI string `json:"verificationUri"`
VerificationURIComplete string `json:"verificationUriComplete"`
ExpiresIn int `json:"expiresIn"`
Interval int `json:"interval"`
}
// CreateTokenResponse from AWS SSO OIDC.
type CreateTokenResponse struct {
AccessToken string `json:"accessToken"`
TokenType string `json:"tokenType"`
ExpiresIn int `json:"expiresIn"`
RefreshToken string `json:"refreshToken"`
}
// RegisterClient registers a new OIDC client with AWS.
func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResponse, error) {
// Generate unique client name for each registration to support multiple accounts
clientName := fmt.Sprintf("CLI-Proxy-API-%d", time.Now().UnixNano())
payload := map[string]interface{}{
"clientName": clientName,
"clientType": "public",
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations"},
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/client/register", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode)
}
var result RegisterClientResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
return &result, nil
}
// StartDeviceAuthorization starts the device authorization flow.
func (c *SSOOIDCClient) StartDeviceAuthorization(ctx context.Context, clientID, clientSecret string) (*StartDeviceAuthResponse, error) {
payload := map[string]string{
"clientId": clientID,
"clientSecret": clientSecret,
"startUrl": builderIDStartURL,
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/device_authorization", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode)
}
var result StartDeviceAuthResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
return &result, nil
}
// CreateToken polls for the access token after user authorization.
func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret, deviceCode string) (*CreateTokenResponse, error) {
payload := map[string]string{
"clientId": clientID,
"clientSecret": clientSecret,
"deviceCode": deviceCode,
"grantType": "urn:ietf:params:oauth:grant-type:device_code",
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
// Check for pending authorization
if resp.StatusCode == http.StatusBadRequest {
var errResp struct {
Error string `json:"error"`
}
if json.Unmarshal(respBody, &errResp) == nil {
if errResp.Error == "authorization_pending" {
return nil, fmt.Errorf("authorization_pending")
}
if errResp.Error == "slow_down" {
return nil, fmt.Errorf("slow_down")
}
}
log.Debugf("create token failed: %s", string(respBody))
return nil, fmt.Errorf("create token failed")
}
if resp.StatusCode != http.StatusOK {
log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode)
}
var result CreateTokenResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
return &result, nil
}
// RefreshToken refreshes an access token using the refresh token.
func (c *SSOOIDCClient) RefreshToken(ctx context.Context, clientID, clientSecret, refreshToken string) (*KiroTokenData, error) {
payload := map[string]string{
"clientId": clientID,
"clientSecret": clientSecret,
"refreshToken": refreshToken,
"grantType": "refresh_token",
}
body, err := json.Marshal(payload)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ssoOIDCEndpoint+"/token", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
log.Debugf("token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
}
var result CreateTokenResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, err
}
expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second)
return &KiroTokenData{
AccessToken: result.AccessToken,
RefreshToken: result.RefreshToken,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "builder-id",
Provider: "AWS",
ClientID: clientID,
ClientSecret: clientSecret,
}, nil
}
// LoginWithBuilderID performs the full device code flow for AWS Builder ID.
func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData, error) {
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
fmt.Println("║ Kiro Authentication (AWS Builder ID) ║")
fmt.Println("╚══════════════════════════════════════════════════════════╝")
// Step 1: Register client
fmt.Println("\nRegistering client...")
regResp, err := c.RegisterClient(ctx)
if err != nil {
return nil, fmt.Errorf("failed to register client: %w", err)
}
log.Debugf("Client registered: %s", regResp.ClientID)
// Step 2: Start device authorization
fmt.Println("Starting device authorization...")
authResp, err := c.StartDeviceAuthorization(ctx, regResp.ClientID, regResp.ClientSecret)
if err != nil {
return nil, fmt.Errorf("failed to start device auth: %w", err)
}
// Step 3: Show user the verification URL
fmt.Printf("\n")
fmt.Println("════════════════════════════════════════════════════════════")
fmt.Printf(" Open this URL in your browser:\n")
fmt.Printf(" %s\n", authResp.VerificationURIComplete)
fmt.Println("════════════════════════════════════════════════════════════")
fmt.Printf("\n Or go to: %s\n", authResp.VerificationURI)
fmt.Printf(" And enter code: %s\n\n", authResp.UserCode)
// Set incognito mode based on config (defaults to true for Kiro, can be overridden with --no-incognito)
// Incognito mode enables multi-account support by bypassing cached sessions
if c.cfg != nil {
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
if !c.cfg.IncognitoBrowser {
log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
} else {
log.Debug("kiro: using incognito mode for multi-account support")
}
} else {
browser.SetIncognitoMode(true) // Default to incognito if no config
log.Debug("kiro: using incognito mode for multi-account support (default)")
}
// Open browser using cross-platform browser package
if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil {
log.Warnf("Could not open browser automatically: %v", err)
fmt.Println(" Please open the URL manually in your browser.")
} else {
fmt.Println(" (Browser opened automatically)")
}
// Step 4: Poll for token
fmt.Println("Waiting for authorization...")
interval := pollInterval
if authResp.Interval > 0 {
interval = time.Duration(authResp.Interval) * time.Second
}
deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second)
for time.Now().Before(deadline) {
select {
case <-ctx.Done():
browser.CloseBrowser() // Cleanup on cancel
return nil, ctx.Err()
case <-time.After(interval):
tokenResp, err := c.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode)
if err != nil {
errStr := err.Error()
if strings.Contains(errStr, "authorization_pending") {
fmt.Print(".")
continue
}
if strings.Contains(errStr, "slow_down") {
interval += 5 * time.Second
continue
}
// Close browser on error before returning
browser.CloseBrowser()
return nil, fmt.Errorf("token creation failed: %w", err)
}
fmt.Println("\n\n✓ Authorization successful!")
// Close the browser window
if err := browser.CloseBrowser(); err != nil {
log.Debugf("Failed to close browser: %v", err)
}
// Step 5: Get profile ARN from CodeWhisperer API
fmt.Println("Fetching profile information...")
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
// Extract email from JWT access token
email := ExtractEmailFromJWT(tokenResp.AccessToken)
if email != "" {
fmt.Printf(" Logged in as: %s\n", email)
}
expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
return &KiroTokenData{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
ProfileArn: profileArn,
ExpiresAt: expiresAt.Format(time.RFC3339),
AuthMethod: "builder-id",
Provider: "AWS",
ClientID: regResp.ClientID,
ClientSecret: regResp.ClientSecret,
Email: email,
}, nil
}
}
// Close browser on timeout for better UX
if err := browser.CloseBrowser(); err != nil {
log.Debugf("Failed to close browser on timeout: %v", err)
}
return nil, fmt.Errorf("authorization timed out")
}
// fetchProfileArn retrieves the profile ARN from CodeWhisperer API.
// This is needed for file naming since AWS SSO OIDC doesn't return profile info.
func (c *SSOOIDCClient) fetchProfileArn(ctx context.Context, accessToken string) string {
// Try ListProfiles API first
profileArn := c.tryListProfiles(ctx, accessToken)
if profileArn != "" {
return profileArn
}
// Fallback: Try ListAvailableCustomizations
return c.tryListCustomizations(ctx, accessToken)
}
func (c *SSOOIDCClient) tryListProfiles(ctx context.Context, accessToken string) string {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
}
body, err := json.Marshal(payload)
if err != nil {
return ""
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
if err != nil {
return ""
}
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListProfiles")
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return ""
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
log.Debugf("ListProfiles failed (status %d): %s", resp.StatusCode, string(respBody))
return ""
}
log.Debugf("ListProfiles response: %s", string(respBody))
var result struct {
Profiles []struct {
Arn string `json:"arn"`
} `json:"profiles"`
ProfileArn string `json:"profileArn"`
}
if err := json.Unmarshal(respBody, &result); err != nil {
return ""
}
if result.ProfileArn != "" {
return result.ProfileArn
}
if len(result.Profiles) > 0 {
return result.Profiles[0].Arn
}
return ""
}
func (c *SSOOIDCClient) tryListCustomizations(ctx context.Context, accessToken string) string {
payload := map[string]interface{}{
"origin": "AI_EDITOR",
}
body, err := json.Marshal(payload)
if err != nil {
return ""
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://codewhisperer.us-east-1.amazonaws.com", strings.NewReader(string(body)))
if err != nil {
return ""
}
req.Header.Set("Content-Type", "application/x-amz-json-1.0")
req.Header.Set("x-amz-target", "AmazonCodeWhispererService.ListAvailableCustomizations")
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return ""
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
log.Debugf("ListAvailableCustomizations failed (status %d): %s", resp.StatusCode, string(respBody))
return ""
}
log.Debugf("ListAvailableCustomizations response: %s", string(respBody))
var result struct {
Customizations []struct {
Arn string `json:"arn"`
} `json:"customizations"`
ProfileArn string `json:"profileArn"`
}
if err := json.Unmarshal(respBody, &result); err != nil {
return ""
}
if result.ProfileArn != "" {
return result.ProfileArn
}
if len(result.Customizations) > 0 {
return result.Customizations[0].Arn
}
return ""
}

View File

@@ -0,0 +1,72 @@
package kiro
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
)
// KiroTokenStorage holds the persistent token data for Kiro authentication.
type KiroTokenStorage struct {
// AccessToken is the OAuth2 access token for API access
AccessToken string `json:"access_token"`
// RefreshToken is used to obtain new access tokens
RefreshToken string `json:"refresh_token"`
// ProfileArn is the AWS CodeWhisperer profile ARN
ProfileArn string `json:"profile_arn"`
// ExpiresAt is the timestamp when the token expires
ExpiresAt string `json:"expires_at"`
// AuthMethod indicates the authentication method used
AuthMethod string `json:"auth_method"`
// Provider indicates the OAuth provider
Provider string `json:"provider"`
// LastRefresh is the timestamp of the last token refresh
LastRefresh string `json:"last_refresh"`
}
// SaveTokenToFile persists the token storage to the specified file path.
func (s *KiroTokenStorage) SaveTokenToFile(authFilePath string) error {
dir := filepath.Dir(authFilePath)
if err := os.MkdirAll(dir, 0700); err != nil {
return fmt.Errorf("failed to create directory: %w", err)
}
data, err := json.MarshalIndent(s, "", " ")
if err != nil {
return fmt.Errorf("failed to marshal token storage: %w", err)
}
if err := os.WriteFile(authFilePath, data, 0600); err != nil {
return fmt.Errorf("failed to write token file: %w", err)
}
return nil
}
// LoadFromFile loads token storage from the specified file path.
func LoadFromFile(authFilePath string) (*KiroTokenStorage, error) {
data, err := os.ReadFile(authFilePath)
if err != nil {
return nil, fmt.Errorf("failed to read token file: %w", err)
}
var storage KiroTokenStorage
if err := json.Unmarshal(data, &storage); err != nil {
return nil, fmt.Errorf("failed to parse token file: %w", err)
}
return &storage, nil
}
// ToTokenData converts storage to KiroTokenData for API use.
func (s *KiroTokenStorage) ToTokenData() *KiroTokenData {
return &KiroTokenData{
AccessToken: s.AccessToken,
RefreshToken: s.RefreshToken,
ProfileArn: s.ProfileArn,
ExpiresAt: s.ExpiresAt,
AuthMethod: s.AuthMethod,
Provider: s.Provider,
}
}

View File

@@ -6,14 +6,48 @@ import (
"fmt"
"os/exec"
"runtime"
"sync"
pkgbrowser "github.com/pkg/browser"
log "github.com/sirupsen/logrus"
"github.com/skratchdot/open-golang/open"
)
// incognitoMode controls whether to open URLs in incognito/private mode.
// This is useful for OAuth flows where you want to use a different account.
var incognitoMode bool
// lastBrowserProcess stores the last opened browser process for cleanup
var lastBrowserProcess *exec.Cmd
var browserMutex sync.Mutex
// SetIncognitoMode enables or disables incognito/private browsing mode.
func SetIncognitoMode(enabled bool) {
incognitoMode = enabled
}
// IsIncognitoMode returns whether incognito mode is enabled.
func IsIncognitoMode() bool {
return incognitoMode
}
// CloseBrowser closes the last opened browser process.
func CloseBrowser() error {
browserMutex.Lock()
defer browserMutex.Unlock()
if lastBrowserProcess == nil || lastBrowserProcess.Process == nil {
return nil
}
err := lastBrowserProcess.Process.Kill()
lastBrowserProcess = nil
return err
}
// OpenURL opens the specified URL in the default web browser.
// It first attempts to use a platform-agnostic library and falls back to
// platform-specific commands if that fails.
// It uses the pkg/browser library which provides robust cross-platform support
// for Windows, macOS, and Linux.
// If incognito mode is enabled, it will open in a private/incognito window.
//
// Parameters:
// - url: The URL to open.
@@ -21,16 +55,22 @@ import (
// Returns:
// - An error if the URL cannot be opened, otherwise nil.
func OpenURL(url string) error {
fmt.Printf("Attempting to open URL in browser: %s\n", url)
log.Debugf("Opening URL in browser: %s (incognito=%v)", url, incognitoMode)
// Try using the open-golang library first
err := open.Run(url)
// If incognito mode is enabled, use platform-specific incognito commands
if incognitoMode {
log.Debug("Using incognito mode")
return openURLIncognito(url)
}
// Use pkg/browser for cross-platform support
err := pkgbrowser.OpenURL(url)
if err == nil {
log.Debug("Successfully opened URL using open-golang library")
log.Debug("Successfully opened URL using pkg/browser library")
return nil
}
log.Debugf("open-golang failed: %v, trying platform-specific commands", err)
log.Debugf("pkg/browser failed: %v, trying platform-specific commands", err)
// Fallback to platform-specific commands
return openURLPlatformSpecific(url)
@@ -78,18 +118,394 @@ func openURLPlatformSpecific(url string) error {
return nil
}
// openURLIncognito opens a URL in incognito/private browsing mode.
// It first tries to detect the default browser and use its incognito flag.
// Falls back to a chain of known browsers if detection fails.
//
// Parameters:
// - url: The URL to open.
//
// Returns:
// - An error if the URL cannot be opened, otherwise nil.
func openURLIncognito(url string) error {
// First, try to detect and use the default browser
if cmd := tryDefaultBrowserIncognito(url); cmd != nil {
log.Debugf("Using detected default browser: %s %v", cmd.Path, cmd.Args[1:])
if err := cmd.Start(); err == nil {
storeBrowserProcess(cmd)
log.Debug("Successfully opened URL in default browser's incognito mode")
return nil
}
log.Debugf("Failed to start default browser, trying fallback chain")
}
// Fallback to known browser chain
cmd := tryFallbackBrowsersIncognito(url)
if cmd == nil {
log.Warn("No browser with incognito support found, falling back to normal mode")
return openURLPlatformSpecific(url)
}
log.Debugf("Running incognito command: %s %v", cmd.Path, cmd.Args[1:])
err := cmd.Start()
if err != nil {
log.Warnf("Failed to open incognito browser: %v, falling back to normal mode", err)
return openURLPlatformSpecific(url)
}
storeBrowserProcess(cmd)
log.Debug("Successfully opened URL in incognito/private mode")
return nil
}
// storeBrowserProcess safely stores the browser process for later cleanup.
func storeBrowserProcess(cmd *exec.Cmd) {
browserMutex.Lock()
lastBrowserProcess = cmd
browserMutex.Unlock()
}
// tryDefaultBrowserIncognito attempts to detect the default browser and return
// an exec.Cmd configured with the appropriate incognito flag.
func tryDefaultBrowserIncognito(url string) *exec.Cmd {
switch runtime.GOOS {
case "darwin":
return tryDefaultBrowserMacOS(url)
case "windows":
return tryDefaultBrowserWindows(url)
case "linux":
return tryDefaultBrowserLinux(url)
}
return nil
}
// tryDefaultBrowserMacOS detects the default browser on macOS.
func tryDefaultBrowserMacOS(url string) *exec.Cmd {
// Try to get default browser from Launch Services
out, err := exec.Command("defaults", "read", "com.apple.LaunchServices/com.apple.launchservices.secure", "LSHandlers").Output()
if err != nil {
return nil
}
output := string(out)
var browserName string
// Parse the output to find the http/https handler
if containsBrowserID(output, "com.google.chrome") {
browserName = "chrome"
} else if containsBrowserID(output, "org.mozilla.firefox") {
browserName = "firefox"
} else if containsBrowserID(output, "com.apple.safari") {
browserName = "safari"
} else if containsBrowserID(output, "com.brave.browser") {
browserName = "brave"
} else if containsBrowserID(output, "com.microsoft.edgemac") {
browserName = "edge"
}
return createMacOSIncognitoCmd(browserName, url)
}
// containsBrowserID checks if the LaunchServices output contains a browser ID.
func containsBrowserID(output, bundleID string) bool {
return stringContains(output, bundleID)
}
// stringContains is a simple string contains check.
func stringContains(s, substr string) bool {
return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
(len(s) > 0 && len(substr) > 0 && findSubstring(s, substr)))
}
func findSubstring(s, substr string) bool {
for i := 0; i <= len(s)-len(substr); i++ {
if s[i:i+len(substr)] == substr {
return true
}
}
return false
}
// createMacOSIncognitoCmd creates the appropriate incognito command for macOS browsers.
func createMacOSIncognitoCmd(browserName, url string) *exec.Cmd {
switch browserName {
case "chrome":
// Try direct path first
chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
if _, err := exec.LookPath(chromePath); err == nil {
return exec.Command(chromePath, "--incognito", url)
}
return exec.Command("open", "-na", "Google Chrome", "--args", "--incognito", url)
case "firefox":
return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
case "safari":
// Safari doesn't have CLI incognito, try AppleScript
return tryAppleScriptSafariPrivate(url)
case "brave":
return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
case "edge":
return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
}
return nil
}
// tryAppleScriptSafariPrivate attempts to open Safari in private browsing mode using AppleScript.
func tryAppleScriptSafariPrivate(url string) *exec.Cmd {
// AppleScript to open a new private window in Safari
script := fmt.Sprintf(`
tell application "Safari"
activate
tell application "System Events"
keystroke "n" using {command down, shift down}
delay 0.5
end tell
set URL of document 1 to "%s"
end tell
`, url)
cmd := exec.Command("osascript", "-e", script)
// Test if this approach works by checking if Safari is available
if _, err := exec.LookPath("/Applications/Safari.app/Contents/MacOS/Safari"); err != nil {
log.Debug("Safari not found, AppleScript private window not available")
return nil
}
log.Debug("Attempting Safari private window via AppleScript")
return cmd
}
// tryDefaultBrowserWindows detects the default browser on Windows via registry.
func tryDefaultBrowserWindows(url string) *exec.Cmd {
// Query registry for default browser
out, err := exec.Command("reg", "query",
`HKEY_CURRENT_USER\Software\Microsoft\Windows\Shell\Associations\UrlAssociations\http\UserChoice`,
"/v", "ProgId").Output()
if err != nil {
return nil
}
output := string(out)
var browserName string
// Map ProgId to browser name
if stringContains(output, "ChromeHTML") {
browserName = "chrome"
} else if stringContains(output, "FirefoxURL") {
browserName = "firefox"
} else if stringContains(output, "MSEdgeHTM") {
browserName = "edge"
} else if stringContains(output, "BraveHTML") {
browserName = "brave"
}
return createWindowsIncognitoCmd(browserName, url)
}
// createWindowsIncognitoCmd creates the appropriate incognito command for Windows browsers.
func createWindowsIncognitoCmd(browserName, url string) *exec.Cmd {
switch browserName {
case "chrome":
paths := []string{
"chrome",
`C:\Program Files\Google\Chrome\Application\chrome.exe`,
`C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
}
for _, p := range paths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--incognito", url)
}
}
case "firefox":
if path, err := exec.LookPath("firefox"); err == nil {
return exec.Command(path, "--private-window", url)
}
case "edge":
paths := []string{
"msedge",
`C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
`C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
}
for _, p := range paths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--inprivate", url)
}
}
case "brave":
paths := []string{
`C:\Program Files\BraveSoftware\Brave-Browser\Application\brave.exe`,
`C:\Program Files (x86)\BraveSoftware\Brave-Browser\Application\brave.exe`,
}
for _, p := range paths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--incognito", url)
}
}
}
return nil
}
// tryDefaultBrowserLinux detects the default browser on Linux using xdg-settings.
func tryDefaultBrowserLinux(url string) *exec.Cmd {
out, err := exec.Command("xdg-settings", "get", "default-web-browser").Output()
if err != nil {
return nil
}
desktop := string(out)
var browserName string
// Map .desktop file to browser name
if stringContains(desktop, "google-chrome") || stringContains(desktop, "chrome") {
browserName = "chrome"
} else if stringContains(desktop, "firefox") {
browserName = "firefox"
} else if stringContains(desktop, "chromium") {
browserName = "chromium"
} else if stringContains(desktop, "brave") {
browserName = "brave"
} else if stringContains(desktop, "microsoft-edge") || stringContains(desktop, "msedge") {
browserName = "edge"
}
return createLinuxIncognitoCmd(browserName, url)
}
// createLinuxIncognitoCmd creates the appropriate incognito command for Linux browsers.
func createLinuxIncognitoCmd(browserName, url string) *exec.Cmd {
switch browserName {
case "chrome":
paths := []string{"google-chrome", "google-chrome-stable"}
for _, p := range paths {
if path, err := exec.LookPath(p); err == nil {
return exec.Command(path, "--incognito", url)
}
}
case "firefox":
paths := []string{"firefox", "firefox-esr"}
for _, p := range paths {
if path, err := exec.LookPath(p); err == nil {
return exec.Command(path, "--private-window", url)
}
}
case "chromium":
paths := []string{"chromium", "chromium-browser"}
for _, p := range paths {
if path, err := exec.LookPath(p); err == nil {
return exec.Command(path, "--incognito", url)
}
}
case "brave":
if path, err := exec.LookPath("brave-browser"); err == nil {
return exec.Command(path, "--incognito", url)
}
case "edge":
if path, err := exec.LookPath("microsoft-edge"); err == nil {
return exec.Command(path, "--inprivate", url)
}
}
return nil
}
// tryFallbackBrowsersIncognito tries a chain of known browsers as fallback.
func tryFallbackBrowsersIncognito(url string) *exec.Cmd {
switch runtime.GOOS {
case "darwin":
return tryFallbackBrowsersMacOS(url)
case "windows":
return tryFallbackBrowsersWindows(url)
case "linux":
return tryFallbackBrowsersLinuxChain(url)
}
return nil
}
// tryFallbackBrowsersMacOS tries known browsers on macOS.
func tryFallbackBrowsersMacOS(url string) *exec.Cmd {
// Try Chrome
chromePath := "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
if _, err := exec.LookPath(chromePath); err == nil {
return exec.Command(chromePath, "--incognito", url)
}
// Try Firefox
if _, err := exec.LookPath("/Applications/Firefox.app/Contents/MacOS/firefox"); err == nil {
return exec.Command("open", "-na", "Firefox", "--args", "--private-window", url)
}
// Try Brave
if _, err := exec.LookPath("/Applications/Brave Browser.app/Contents/MacOS/Brave Browser"); err == nil {
return exec.Command("open", "-na", "Brave Browser", "--args", "--incognito", url)
}
// Try Edge
if _, err := exec.LookPath("/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge"); err == nil {
return exec.Command("open", "-na", "Microsoft Edge", "--args", "--inprivate", url)
}
// Last resort: try Safari with AppleScript
if cmd := tryAppleScriptSafariPrivate(url); cmd != nil {
log.Info("Using Safari with AppleScript for private browsing (may require accessibility permissions)")
return cmd
}
return nil
}
// tryFallbackBrowsersWindows tries known browsers on Windows.
func tryFallbackBrowsersWindows(url string) *exec.Cmd {
// Chrome
chromePaths := []string{
"chrome",
`C:\Program Files\Google\Chrome\Application\chrome.exe`,
`C:\Program Files (x86)\Google\Chrome\Application\chrome.exe`,
}
for _, p := range chromePaths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--incognito", url)
}
}
// Firefox
if path, err := exec.LookPath("firefox"); err == nil {
return exec.Command(path, "--private-window", url)
}
// Edge (usually available on Windows 10+)
edgePaths := []string{
"msedge",
`C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe`,
`C:\Program Files\Microsoft\Edge\Application\msedge.exe`,
}
for _, p := range edgePaths {
if _, err := exec.LookPath(p); err == nil {
return exec.Command(p, "--inprivate", url)
}
}
return nil
}
// tryFallbackBrowsersLinuxChain tries known browsers on Linux.
func tryFallbackBrowsersLinuxChain(url string) *exec.Cmd {
type browserConfig struct {
name string
flag string
}
browsers := []browserConfig{
{"google-chrome", "--incognito"},
{"google-chrome-stable", "--incognito"},
{"chromium", "--incognito"},
{"chromium-browser", "--incognito"},
{"firefox", "--private-window"},
{"firefox-esr", "--private-window"},
{"brave-browser", "--incognito"},
{"microsoft-edge", "--inprivate"},
}
for _, b := range browsers {
if path, err := exec.LookPath(b.name); err == nil {
return exec.Command(path, b.flag, url)
}
}
return nil
}
// IsAvailable checks if the system has a command available to open a web browser.
// It verifies the presence of necessary commands for the current operating system.
//
// Returns:
// - true if a browser can be opened, false otherwise.
func IsAvailable() bool {
// First check if open-golang can work
testErr := open.Run("about:blank")
if testErr == nil {
return true
}
// Check platform-specific commands
switch runtime.GOOS {
case "darwin":

View File

@@ -19,6 +19,7 @@ func newAuthManager() *sdkAuth.Manager {
sdkAuth.NewQwenAuthenticator(),
sdkAuth.NewIFlowAuthenticator(),
sdkAuth.NewAntigravityAuthenticator(),
sdkAuth.NewKiroAuthenticator(),
)
return manager
}

160
internal/cmd/kiro_login.go Normal file
View File

@@ -0,0 +1,160 @@
package cmd
import (
"context"
"fmt"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
log "github.com/sirupsen/logrus"
)
// DoKiroLogin triggers the Kiro authentication flow with Google OAuth.
// This is the default login method (same as --kiro-google-login).
//
// Parameters:
// - cfg: The application configuration
// - options: Login options including Prompt field
func DoKiroLogin(cfg *config.Config, options *LoginOptions) {
// Use Google login as default
DoKiroGoogleLogin(cfg, options)
}
// DoKiroGoogleLogin triggers Kiro authentication with Google OAuth.
// This uses a custom protocol handler (kiro://) to receive the callback.
//
// Parameters:
// - cfg: The application configuration
// - options: Login options including prompts
func DoKiroGoogleLogin(cfg *config.Config, options *LoginOptions) {
if options == nil {
options = &LoginOptions{}
}
// Note: Kiro defaults to incognito mode for multi-account support.
// Users can override with --no-incognito if they want to use existing browser sessions.
manager := newAuthManager()
// Use KiroAuthenticator with Google login
authenticator := sdkAuth.NewKiroAuthenticator()
record, err := authenticator.LoginWithGoogle(context.Background(), cfg, &sdkAuth.LoginOptions{
NoBrowser: options.NoBrowser,
Metadata: map[string]string{},
Prompt: options.Prompt,
})
if err != nil {
log.Errorf("Kiro Google authentication failed: %v", err)
fmt.Println("\nTroubleshooting:")
fmt.Println("1. Make sure the protocol handler is installed")
fmt.Println("2. Complete the Google login in the browser")
fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
return
}
// Save the auth record
savedPath, err := manager.SaveAuth(record, cfg)
if err != nil {
log.Errorf("Failed to save auth: %v", err)
return
}
if savedPath != "" {
fmt.Printf("Authentication saved to %s\n", savedPath)
}
if record != nil && record.Label != "" {
fmt.Printf("Authenticated as %s\n", record.Label)
}
fmt.Println("Kiro Google authentication successful!")
}
// DoKiroAWSLogin triggers Kiro authentication with AWS Builder ID.
// This uses the device code flow for AWS SSO OIDC authentication.
//
// Parameters:
// - cfg: The application configuration
// - options: Login options including prompts
func DoKiroAWSLogin(cfg *config.Config, options *LoginOptions) {
if options == nil {
options = &LoginOptions{}
}
// Note: Kiro defaults to incognito mode for multi-account support.
// Users can override with --no-incognito if they want to use existing browser sessions.
manager := newAuthManager()
// Use KiroAuthenticator with AWS Builder ID login (device code flow)
authenticator := sdkAuth.NewKiroAuthenticator()
record, err := authenticator.Login(context.Background(), cfg, &sdkAuth.LoginOptions{
NoBrowser: options.NoBrowser,
Metadata: map[string]string{},
Prompt: options.Prompt,
})
if err != nil {
log.Errorf("Kiro AWS authentication failed: %v", err)
fmt.Println("\nTroubleshooting:")
fmt.Println("1. Make sure you have an AWS Builder ID")
fmt.Println("2. Complete the authorization in the browser")
fmt.Println("3. If callback fails, try: --kiro-import (after logging in via Kiro IDE)")
return
}
// Save the auth record
savedPath, err := manager.SaveAuth(record, cfg)
if err != nil {
log.Errorf("Failed to save auth: %v", err)
return
}
if savedPath != "" {
fmt.Printf("Authentication saved to %s\n", savedPath)
}
if record != nil && record.Label != "" {
fmt.Printf("Authenticated as %s\n", record.Label)
}
fmt.Println("Kiro AWS authentication successful!")
}
// DoKiroImport imports Kiro token from Kiro IDE's token file.
// This is useful for users who have already logged in via Kiro IDE
// and want to use the same credentials in CLI Proxy API.
//
// Parameters:
// - cfg: The application configuration
// - options: Login options (currently unused for import)
func DoKiroImport(cfg *config.Config, options *LoginOptions) {
if options == nil {
options = &LoginOptions{}
}
manager := newAuthManager()
// Use ImportFromKiroIDE instead of Login
authenticator := sdkAuth.NewKiroAuthenticator()
record, err := authenticator.ImportFromKiroIDE(context.Background(), cfg)
if err != nil {
log.Errorf("Kiro token import failed: %v", err)
fmt.Println("\nMake sure you have logged in to Kiro IDE first:")
fmt.Println("1. Open Kiro IDE")
fmt.Println("2. Click 'Sign in with Google' (or GitHub)")
fmt.Println("3. Complete the login process")
fmt.Println("4. Run this command again")
return
}
// Save the imported auth record
savedPath, err := manager.SaveAuth(record, cfg)
if err != nil {
log.Errorf("Failed to save auth: %v", err)
return
}
if savedPath != "" {
fmt.Printf("Authentication saved to %s\n", savedPath)
}
if record != nil && record.Label != "" {
fmt.Printf("Imported as %s\n", record.Label)
}
fmt.Println("Kiro token import successful!")
}

View File

@@ -58,6 +58,9 @@ type Config struct {
// GeminiKey defines Gemini API key configurations with optional routing overrides.
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
// KiroKey defines a list of Kiro (AWS CodeWhisperer) configurations.
KiroKey []KiroKey `yaml:"kiro" json:"kiro"`
// Codex defines a list of Codex API key configurations as specified in the YAML configuration file.
CodexKey []CodexKey `yaml:"codex-api-key" json:"codex-api-key"`
@@ -80,6 +83,11 @@ type Config struct {
// Payload defines default and override rules for provider payload parameters.
Payload PayloadConfig `yaml:"payload" json:"payload"`
// IncognitoBrowser enables opening OAuth URLs in incognito/private browsing mode.
// This is useful when you want to login with a different account without logging out
// from your current session. Default: false.
IncognitoBrowser bool `yaml:"incognito-browser" json:"incognito-browser"`
legacyMigrationPending bool `yaml:"-" json:"-"`
}
@@ -240,6 +248,31 @@ type GeminiKey struct {
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
}
// KiroKey represents the configuration for Kiro (AWS CodeWhisperer) authentication.
type KiroKey struct {
// TokenFile is the path to the Kiro token file (default: ~/.aws/sso/cache/kiro-auth-token.json)
TokenFile string `yaml:"token-file,omitempty" json:"token-file,omitempty"`
// AccessToken is the OAuth access token for direct configuration.
AccessToken string `yaml:"access-token,omitempty" json:"access-token,omitempty"`
// RefreshToken is the OAuth refresh token for token renewal.
RefreshToken string `yaml:"refresh-token,omitempty" json:"refresh-token,omitempty"`
// ProfileArn is the AWS CodeWhisperer profile ARN.
ProfileArn string `yaml:"profile-arn,omitempty" json:"profile-arn,omitempty"`
// Region is the AWS region (default: us-east-1).
Region string `yaml:"region,omitempty" json:"region,omitempty"`
// ProxyURL optionally overrides the global proxy for this configuration.
ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
// AgentTaskType sets the Kiro API task type. Known values: "vibe", "dev", "chat".
// Leave empty to let API use defaults. Different values may inject different system prompts.
AgentTaskType string `yaml:"agent-task-type,omitempty" json:"agent-task-type,omitempty"`
}
// OpenAICompatibility represents the configuration for OpenAI API compatibility
// with external providers, allowing model aliases to be routed through OpenAI API format.
type OpenAICompatibility struct {
@@ -320,6 +353,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
cfg.UsageStatisticsEnabled = false
cfg.DisableCooling = false
cfg.AmpCode.RestrictManagementToLocalhost = true // Default to secure: only localhost access
cfg.IncognitoBrowser = false // Default to normal browser (AWS uses incognito by force)
if err = yaml.Unmarshal(data, &cfg); err != nil {
if optional {
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
@@ -370,6 +404,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
// Sanitize Claude key headers
cfg.SanitizeClaudeKeys()
// Sanitize Kiro keys: trim whitespace from credential fields
cfg.SanitizeKiroKeys()
// Sanitize OpenAI compatibility providers: drop entries without base-url
cfg.SanitizeOpenAICompatibility()
@@ -446,6 +483,22 @@ func (cfg *Config) SanitizeClaudeKeys() {
}
}
// SanitizeKiroKeys trims whitespace from Kiro credential fields.
func (cfg *Config) SanitizeKiroKeys() {
if cfg == nil || len(cfg.KiroKey) == 0 {
return
}
for i := range cfg.KiroKey {
entry := &cfg.KiroKey[i]
entry.TokenFile = strings.TrimSpace(entry.TokenFile)
entry.AccessToken = strings.TrimSpace(entry.AccessToken)
entry.RefreshToken = strings.TrimSpace(entry.RefreshToken)
entry.ProfileArn = strings.TrimSpace(entry.ProfileArn)
entry.Region = strings.TrimSpace(entry.Region)
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
}
}
// SanitizeGeminiKeys deduplicates and normalizes Gemini credentials.
func (cfg *Config) SanitizeGeminiKeys() {
if cfg == nil {

View File

@@ -24,4 +24,7 @@ const (
// Antigravity represents the Antigravity response format identifier.
Antigravity = "antigravity"
// Kiro represents the AWS CodeWhisperer (Kiro) provider identifier.
Kiro = "kiro"
)

View File

@@ -38,13 +38,16 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
timestamp := entry.Time.Format("2006-01-02 15:04:05")
message := strings.TrimRight(entry.Message, "\r\n")
var formatted string
// Handle nil Caller (can happen with some log entries)
callerFile := "unknown"
callerLine := 0
if entry.Caller != nil {
formatted = fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, filepath.Base(entry.Caller.File), entry.Caller.Line, message)
} else {
formatted = fmt.Sprintf("[%s] [%s] %s\n", timestamp, entry.Level, message)
callerFile = filepath.Base(entry.Caller.File)
callerLine = entry.Caller.Line
}
formatted := fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, callerFile, callerLine, message)
buffer.WriteString(formatted)
return buffer.Bytes(), nil
@@ -55,6 +58,7 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
func SetupBaseLogger() {
setupOnce.Do(func() {
log.SetOutput(os.Stdout)
log.SetLevel(log.InfoLevel)
log.SetReportCaller(true)
log.SetFormatter(&LogFormatter{})

View File

@@ -986,3 +986,161 @@ func GetIFlowModels() []*ModelInfo {
}
return models
}
// GetKiroModels returns the Kiro (AWS CodeWhisperer) model definitions
func GetKiroModels() []*ModelInfo {
return []*ModelInfo{
{
ID: "kiro-auto",
Object: "model",
Created: 1732752000, // 2024-11-28
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Auto",
Description: "Automatic model selection by AWS CodeWhisperer",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-opus-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Opus 4.5",
Description: "Claude Opus 4.5 via Kiro (2.2x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-sonnet-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Sonnet 4.5",
Description: "Claude Sonnet 4.5 via Kiro (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-sonnet-4",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Sonnet 4",
Description: "Claude Sonnet 4 via Kiro (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-haiku-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Haiku 4.5",
Description: "Claude Haiku 4.5 via Kiro (0.4x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
// --- Chat Variant (No tool calling, for pure conversation) ---
{
ID: "kiro-claude-opus-4.5-chat",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Opus 4.5 (Chat)",
Description: "Claude Opus 4.5 for chat only (no tool calling)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
// --- Agentic Variants (Optimized for coding agents with chunked writes) ---
{
ID: "kiro-claude-opus-4.5-agentic",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Opus 4.5 (Agentic)",
Description: "Claude Opus 4.5 optimized for coding agents (chunked writes)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "kiro-claude-sonnet-4.5-agentic",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Kiro Claude Sonnet 4.5 (Agentic)",
Description: "Claude Sonnet 4.5 optimized for coding agents (chunked writes)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
}
}
// GetAmazonQModels returns the Amazon Q (AWS CodeWhisperer) model definitions.
// These models use the same API as Kiro and share the same executor.
func GetAmazonQModels() []*ModelInfo {
return []*ModelInfo{
{
ID: "amazonq-auto",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro", // Uses Kiro executor - same API
DisplayName: "Amazon Q Auto",
Description: "Automatic model selection by Amazon Q",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-opus-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Opus 4.5",
Description: "Claude Opus 4.5 via Amazon Q (2.2x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-sonnet-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Sonnet 4.5",
Description: "Claude Sonnet 4.5 via Amazon Q (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-sonnet-4",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Sonnet 4",
Description: "Claude Sonnet 4 via Amazon Q (1.3x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
{
ID: "amazonq-claude-haiku-4.5",
Object: "model",
Created: 1732752000,
OwnedBy: "aws",
Type: "kiro",
DisplayName: "Amazon Q Claude Haiku 4.5",
Description: "Claude Haiku 4.5 via Amazon Q (0.4x credit)",
ContextLength: 200000,
MaxCompletionTokens: 64000,
},
}
}

View File

@@ -12,6 +12,7 @@ import (
"net/url"
"strconv"
"strings"
"sync"
"time"
"github.com/google/uuid"
@@ -41,7 +42,10 @@ const (
streamScannerBuffer int = 20_971_520
)
var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
var (
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
randSourceMutex sync.Mutex
)
// AntigravityExecutor proxies requests to the antigravity upstream.
type AntigravityExecutor struct {
@@ -754,15 +758,19 @@ func generateRequestID() string {
}
func generateSessionID() string {
randSourceMutex.Lock()
n := randSource.Int63n(9_000_000_000_000_000_000)
randSourceMutex.Unlock()
return "-" + strconv.FormatInt(n, 10)
}
func generateProjectID() string {
adjectives := []string{"useful", "bright", "swift", "calm", "bold"}
nouns := []string{"fuze", "wave", "spark", "flow", "core"}
randSourceMutex.Lock()
adj := adjectives[randSource.Intn(len(adjectives))]
noun := nouns[randSource.Intn(len(nouns))]
randSourceMutex.Unlock()
randomPart := strings.ToLower(uuid.NewString())[:5]
return adj + "-" + noun + "-" + randomPart
}

View File

@@ -1,10 +1,38 @@
package executor
import "time"
import (
"sync"
"time"
)
type codexCache struct {
ID string
Expire time.Time
}
var codexCacheMap = map[string]codexCache{}
var (
codexCacheMap = map[string]codexCache{}
codexCacheMutex sync.RWMutex
)
// getCodexCache safely retrieves a cache entry
func getCodexCache(key string) (codexCache, bool) {
codexCacheMutex.RLock()
defer codexCacheMutex.RUnlock()
cache, ok := codexCacheMap[key]
return cache, ok
}
// setCodexCache safely sets a cache entry
func setCodexCache(key string, cache codexCache) {
codexCacheMutex.Lock()
defer codexCacheMutex.Unlock()
codexCacheMap[key] = cache
}
// deleteCodexCache safely deletes a cache entry
func deleteCodexCache(key string) {
codexCacheMutex.Lock()
defer codexCacheMutex.Unlock()
delete(codexCacheMap, key)
}

View File

@@ -506,12 +506,12 @@ func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Form
if userIDResult.Exists() {
var hasKey bool
key := fmt.Sprintf("%s-%s", req.Model, userIDResult.String())
if cache, hasKey = codexCacheMap[key]; !hasKey || cache.Expire.Before(time.Now()) {
if cache, hasKey = getCodexCache(key); !hasKey || cache.Expire.Before(time.Now()) {
cache = codexCache{
ID: uuid.New().String(),
Expire: time.Now().Add(1 * time.Hour),
}
codexCacheMap[key] = cache
setCodexCache(key, cache)
}
}
} else if from == "openai-response" {

File diff suppressed because it is too large Load Diff

View File

@@ -6,6 +6,7 @@ import (
"net/http"
"net/url"
"strings"
"sync"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
@@ -14,11 +15,19 @@ import (
"golang.org/x/net/proxy"
)
// httpClientCache caches HTTP clients by proxy URL to enable connection reuse
var (
httpClientCache = make(map[string]*http.Client)
httpClientCacheMutex sync.RWMutex
)
// newProxyAwareHTTPClient creates an HTTP client with proper proxy configuration priority:
// 1. Use auth.ProxyURL if configured (highest priority)
// 2. Use cfg.ProxyURL if auth proxy is not configured
// 3. Use RoundTripper from context if neither are configured
//
// This function caches HTTP clients by proxy URL to enable TCP/TLS connection reuse.
//
// Parameters:
// - ctx: The context containing optional RoundTripper
// - cfg: The application configuration
@@ -28,11 +37,6 @@ import (
// Returns:
// - *http.Client: An HTTP client with configured proxy or transport
func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
httpClient := &http.Client{}
if timeout > 0 {
httpClient.Timeout = timeout
}
// Priority 1: Use auth.ProxyURL if configured
var proxyURL string
if auth != nil {
@@ -44,11 +48,39 @@ func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *clip
proxyURL = strings.TrimSpace(cfg.ProxyURL)
}
// Build cache key from proxy URL (empty string for no proxy)
cacheKey := proxyURL
// Check cache first
httpClientCacheMutex.RLock()
if cachedClient, ok := httpClientCache[cacheKey]; ok {
httpClientCacheMutex.RUnlock()
// Return a wrapper with the requested timeout but shared transport
if timeout > 0 {
return &http.Client{
Transport: cachedClient.Transport,
Timeout: timeout,
}
}
return cachedClient
}
httpClientCacheMutex.RUnlock()
// Create new client
httpClient := &http.Client{}
if timeout > 0 {
httpClient.Timeout = timeout
}
// If we have a proxy URL configured, set up the transport
if proxyURL != "" {
transport := buildProxyTransport(proxyURL)
if transport != nil {
httpClient.Transport = transport
// Cache the client
httpClientCacheMutex.Lock()
httpClientCache[cacheKey] = httpClient
httpClientCacheMutex.Unlock()
return httpClient
}
// If proxy setup failed, log and fall through to context RoundTripper
@@ -60,6 +92,13 @@ func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *clip
httpClient.Transport = rt
}
// Cache the client for no-proxy case
if proxyURL == "" {
httpClientCacheMutex.Lock()
httpClientCache[cacheKey] = httpClient
httpClientCacheMutex.Unlock()
}
return httpClient
}

View File

@@ -331,8 +331,9 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
streamingEvents := make([][]byte, 0)
scanner := bufio.NewScanner(bytes.NewReader(rawJSON))
buffer := make([]byte, 20_971_520)
scanner.Buffer(buffer, 20_971_520)
// Use a smaller initial buffer (64KB) that can grow up to 20MB if needed
// This prevents allocating 20MB for every request regardless of size
scanner.Buffer(make([]byte, 64*1024), 20_971_520)
for scanner.Scan() {
line := scanner.Bytes()
// log.Debug(string(line))

View File

@@ -50,6 +50,10 @@ type ToolCallAccumulator struct {
// Returns:
// - []string: A slice of strings, each containing an OpenAI-compatible JSON response
func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
var localParam any
if param == nil {
param = &localParam
}
if *param == nil {
*param = &ConvertAnthropicResponseToOpenAIParams{
CreatedAt: 0,

View File

@@ -60,12 +60,12 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
// Track whether tools are being used in this response chunk
usedTool := false
output := ""
var sb strings.Builder
// Initialize the streaming session with a message_start event
// This is only sent for the very first response chunk to establish the streaming session
if !(*param).(*Params).HasFirstResponse {
output = "event: message_start\n"
sb.WriteString("event: message_start\n")
// Create the initial message structure with default values according to Claude Code API specification
// This follows the Claude Code API specification for streaming message initialization
@@ -78,7 +78,7 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
if responseIDResult := gjson.GetBytes(rawJSON, "response.responseId"); responseIDResult.Exists() {
messageStartTemplate, _ = sjson.Set(messageStartTemplate, "message.id", responseIDResult.String())
}
output = output + fmt.Sprintf("data: %s\n\n\n", messageStartTemplate)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", messageStartTemplate))
(*param).(*Params).HasFirstResponse = true
}
@@ -101,62 +101,52 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
if partResult.Get("thought").Bool() {
// Continue existing thinking block if already in thinking state
if (*param).(*Params).ResponseType == 2 {
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, (*param).(*Params).ResponseIndex), "delta.thinking", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
} else {
// Transition from another state to thinking
// First, close any existing content block
if (*param).(*Params).ResponseType != 0 {
if (*param).(*Params).ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
// output = output + "\n\n\n"
}
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
}
// Start a new thinking content block
output = output + "event: content_block_start\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"thinking","thinking":""}}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_start\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"thinking","thinking":""}}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, (*param).(*Params).ResponseIndex), "delta.thinking", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
(*param).(*Params).ResponseType = 2 // Set state to thinking
}
} else {
// Process regular text content (user-visible output)
// Continue existing text block if already in content state
if (*param).(*Params).ResponseType == 1 {
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, (*param).(*Params).ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
} else {
// Transition from another state to text content
// First, close any existing content block
if (*param).(*Params).ResponseType != 0 {
if (*param).(*Params).ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
// output = output + "\n\n\n"
}
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
}
// Start a new text content block
output = output + "event: content_block_start\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_start\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
sb.WriteString("event: content_block_delta\n")
data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, (*param).(*Params).ResponseIndex), "delta.text", partTextResult.String())
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
(*param).(*Params).ResponseType = 1 // Set state to content
}
}
@@ -169,42 +159,35 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
// Handle state transitions when switching to function calls
// Close any existing function call block first
if (*param).(*Params).ResponseType == 3 {
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
(*param).(*Params).ResponseType = 0
}
// Special handling for thinking state transition
if (*param).(*Params).ResponseType == 2 {
// output = output + "event: content_block_delta\n"
// output = output + fmt.Sprintf(`data: {"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":null}}`, (*param).(*Params).ResponseIndex)
// output = output + "\n\n\n"
}
// Close any other existing content block
if (*param).(*Params).ResponseType != 0 {
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
(*param).(*Params).ResponseIndex++
}
// Start a new tool use content block
// This creates the structure for a function call in Claude Code format
output = output + "event: content_block_start\n"
sb.WriteString("event: content_block_start\n")
// Create the tool use block with unique ID and function details
data := fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"tool_use","id":"","name":"","input":{}}}`, (*param).(*Params).ResponseIndex)
data, _ = sjson.Set(data, "content_block.id", fmt.Sprintf("%s-%d", fcName, time.Now().UnixNano()))
data, _ = sjson.Set(data, "content_block.name", fcName)
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
if fcArgsResult := functionCallResult.Get("args"); fcArgsResult.Exists() {
output = output + "event: content_block_delta\n"
sb.WriteString("event: content_block_delta\n")
data, _ = sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"input_json_delta","partial_json":""}}`, (*param).(*Params).ResponseIndex), "delta.partial_json", fcArgsResult.Raw)
output = output + fmt.Sprintf("data: %s\n\n\n", data)
sb.WriteString(fmt.Sprintf("data: %s\n\n\n", data))
}
(*param).(*Params).ResponseType = 3
}
@@ -216,13 +199,13 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
if usageResult.Exists() && bytes.Contains(rawJSON, []byte(`"finishReason"`)) {
if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() {
// Close the final content block
output = output + "event: content_block_stop\n"
output = output + fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex)
output = output + "\n\n\n"
sb.WriteString("event: content_block_stop\n")
sb.WriteString(fmt.Sprintf(`data: {"type":"content_block_stop","index":%d}`, (*param).(*Params).ResponseIndex))
sb.WriteString("\n\n\n")
// Send the final message delta with usage information and stop reason
output = output + "event: message_delta\n"
output = output + `data: `
sb.WriteString("event: message_delta\n")
sb.WriteString(`data: `)
// Create the message delta template with appropriate stop reason
template := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":0,"output_tokens":0}}`
@@ -236,11 +219,11 @@ func ConvertGeminiCLIResponseToClaude(_ context.Context, _ string, originalReque
template, _ = sjson.Set(template, "usage.output_tokens", candidatesTokenCountResult.Int()+thoughtsTokenCount)
template, _ = sjson.Set(template, "usage.input_tokens", usageResult.Get("promptTokenCount").Int())
output = output + template + "\n\n\n"
sb.WriteString(template + "\n\n\n")
}
}
return []string{output}
return []string{sb.String()}
}
// ConvertGeminiCLIResponseToClaudeNonStream converts a non-streaming Gemini CLI response to a non-streaming Claude response.

View File

@@ -33,4 +33,7 @@ import (
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/gemini"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/chat-completions"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/openai/responses"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/claude"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/kiro/openai/chat-completions"
)

View File

@@ -0,0 +1,19 @@
package claude
import (
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
)
func init() {
translator.Register(
Claude,
Kiro,
ConvertClaudeRequestToKiro,
interfaces.TranslateResponse{
Stream: ConvertKiroResponseToClaude,
NonStream: ConvertKiroResponseToClaudeNonStream,
},
)
}

View File

@@ -0,0 +1,24 @@
// Package claude provides translation between Kiro and Claude formats.
// Since Kiro uses Claude-compatible format internally, translations are mostly pass-through.
package claude
import (
"bytes"
"context"
)
// ConvertClaudeRequestToKiro converts Claude request to Kiro format.
// Since Kiro uses Claude format internally, this is mostly a pass-through.
func ConvertClaudeRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
return bytes.Clone(inputRawJSON)
}
// ConvertKiroResponseToClaude converts Kiro streaming response to Claude format.
func ConvertKiroResponseToClaude(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
return []string{string(rawResponse)}
}
// ConvertKiroResponseToClaudeNonStream converts Kiro non-streaming response to Claude format.
func ConvertKiroResponseToClaudeNonStream(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
return string(rawResponse)
}

View File

@@ -0,0 +1,19 @@
package chat_completions
import (
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/translator"
)
func init() {
translator.Register(
OpenAI,
Kiro,
ConvertOpenAIRequestToKiro,
interfaces.TranslateResponse{
Stream: ConvertKiroResponseToOpenAI,
NonStream: ConvertKiroResponseToOpenAINonStream,
},
)
}

View File

@@ -0,0 +1,258 @@
// Package chat_completions provides request translation from OpenAI to Kiro format.
package chat_completions
import (
"bytes"
"encoding/json"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// ConvertOpenAIRequestToKiro transforms an OpenAI Chat Completions API request into Kiro (Claude) format.
// Kiro uses Claude-compatible format internally, so we primarily pass through to Claude format.
// Supports tool calling: OpenAI tools -> Claude tools, tool_calls -> tool_use, tool messages -> tool_result.
func ConvertOpenAIRequestToKiro(modelName string, inputRawJSON []byte, stream bool) []byte {
rawJSON := bytes.Clone(inputRawJSON)
root := gjson.ParseBytes(rawJSON)
// Build Claude-compatible request
out := `{"model":"","max_tokens":32000,"messages":[]}`
// Set model
out, _ = sjson.Set(out, "model", modelName)
// Copy max_tokens if present
if v := root.Get("max_tokens"); v.Exists() {
out, _ = sjson.Set(out, "max_tokens", v.Int())
}
// Copy temperature if present
if v := root.Get("temperature"); v.Exists() {
out, _ = sjson.Set(out, "temperature", v.Float())
}
// Copy top_p if present
if v := root.Get("top_p"); v.Exists() {
out, _ = sjson.Set(out, "top_p", v.Float())
}
// Convert OpenAI tools to Claude tools format
if tools := root.Get("tools"); tools.Exists() && tools.IsArray() {
claudeTools := make([]interface{}, 0)
for _, tool := range tools.Array() {
if tool.Get("type").String() == "function" {
fn := tool.Get("function")
claudeTool := map[string]interface{}{
"name": fn.Get("name").String(),
"description": fn.Get("description").String(),
}
// Convert parameters to input_schema
if params := fn.Get("parameters"); params.Exists() {
claudeTool["input_schema"] = params.Value()
} else {
claudeTool["input_schema"] = map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{},
}
}
claudeTools = append(claudeTools, claudeTool)
}
}
if len(claudeTools) > 0 {
out, _ = sjson.Set(out, "tools", claudeTools)
}
}
// Process messages
messages := root.Get("messages")
if messages.Exists() && messages.IsArray() {
claudeMessages := make([]interface{}, 0)
var systemPrompt string
// Track pending tool results to merge with next user message
var pendingToolResults []map[string]interface{}
for _, msg := range messages.Array() {
role := msg.Get("role").String()
content := msg.Get("content")
if role == "system" {
// Extract system message
if content.IsArray() {
for _, part := range content.Array() {
if part.Get("type").String() == "text" {
systemPrompt += part.Get("text").String() + "\n"
}
}
} else {
systemPrompt = content.String()
}
continue
}
if role == "tool" {
// OpenAI tool message -> Claude tool_result content block
toolCallID := msg.Get("tool_call_id").String()
toolContent := content.String()
toolResult := map[string]interface{}{
"type": "tool_result",
"tool_use_id": toolCallID,
}
// Handle content - can be string or structured
if content.IsArray() {
contentParts := make([]interface{}, 0)
for _, part := range content.Array() {
if part.Get("type").String() == "text" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": part.Get("text").String(),
})
}
}
toolResult["content"] = contentParts
} else {
toolResult["content"] = toolContent
}
pendingToolResults = append(pendingToolResults, toolResult)
continue
}
claudeMsg := map[string]interface{}{
"role": role,
}
// Handle assistant messages with tool_calls
if role == "assistant" && msg.Get("tool_calls").Exists() {
contentParts := make([]interface{}, 0)
// Add text content if present
if content.Exists() && content.String() != "" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": content.String(),
})
}
// Convert tool_calls to tool_use blocks
for _, toolCall := range msg.Get("tool_calls").Array() {
toolUseID := toolCall.Get("id").String()
fnName := toolCall.Get("function.name").String()
fnArgs := toolCall.Get("function.arguments").String()
// Parse arguments JSON
var argsMap map[string]interface{}
if err := json.Unmarshal([]byte(fnArgs), &argsMap); err != nil {
argsMap = map[string]interface{}{"raw": fnArgs}
}
contentParts = append(contentParts, map[string]interface{}{
"type": "tool_use",
"id": toolUseID,
"name": fnName,
"input": argsMap,
})
}
claudeMsg["content"] = contentParts
claudeMessages = append(claudeMessages, claudeMsg)
continue
}
// Handle user messages - may need to include pending tool results
if role == "user" && len(pendingToolResults) > 0 {
contentParts := make([]interface{}, 0)
// Add pending tool results first
for _, tr := range pendingToolResults {
contentParts = append(contentParts, tr)
}
pendingToolResults = nil
// Add user content
if content.IsArray() {
for _, part := range content.Array() {
partType := part.Get("type").String()
if partType == "text" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": part.Get("text").String(),
})
} else if partType == "image_url" {
contentParts = append(contentParts, map[string]interface{}{
"type": "image",
"source": map[string]interface{}{
"type": "url",
"url": part.Get("image_url.url").String(),
},
})
}
}
} else if content.String() != "" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": content.String(),
})
}
claudeMsg["content"] = contentParts
claudeMessages = append(claudeMessages, claudeMsg)
continue
}
// Handle regular content
if content.IsArray() {
contentParts := make([]interface{}, 0)
for _, part := range content.Array() {
partType := part.Get("type").String()
if partType == "text" {
contentParts = append(contentParts, map[string]interface{}{
"type": "text",
"text": part.Get("text").String(),
})
} else if partType == "image_url" {
contentParts = append(contentParts, map[string]interface{}{
"type": "image",
"source": map[string]interface{}{
"type": "url",
"url": part.Get("image_url.url").String(),
},
})
}
}
claudeMsg["content"] = contentParts
} else {
claudeMsg["content"] = content.String()
}
claudeMessages = append(claudeMessages, claudeMsg)
}
// If there are pending tool results without a following user message,
// create a user message with just the tool results
if len(pendingToolResults) > 0 {
contentParts := make([]interface{}, 0)
for _, tr := range pendingToolResults {
contentParts = append(contentParts, tr)
}
claudeMessages = append(claudeMessages, map[string]interface{}{
"role": "user",
"content": contentParts,
})
}
out, _ = sjson.Set(out, "messages", claudeMessages)
if systemPrompt != "" {
out, _ = sjson.Set(out, "system", systemPrompt)
}
}
// Set stream
out, _ = sjson.Set(out, "stream", stream)
return []byte(out)
}

View File

@@ -0,0 +1,316 @@
// Package chat_completions provides response translation from Kiro to OpenAI format.
package chat_completions
import (
"context"
"encoding/json"
"time"
"github.com/google/uuid"
"github.com/tidwall/gjson"
)
// ConvertKiroResponseToOpenAI converts Kiro streaming response to OpenAI SSE format.
// Handles Claude SSE events: content_block_start, content_block_delta, input_json_delta,
// content_block_stop, message_delta, and message_stop.
func ConvertKiroResponseToOpenAI(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) []string {
root := gjson.ParseBytes(rawResponse)
var results []string
eventType := root.Get("type").String()
switch eventType {
case "message_start":
// Initial message event - could emit initial chunk if needed
return results
case "content_block_start":
// Start of a content block (text or tool_use)
blockType := root.Get("content_block.type").String()
index := int(root.Get("index").Int())
if blockType == "tool_use" {
// Start of tool_use block
toolUseID := root.Get("content_block.id").String()
toolName := root.Get("content_block.name").String()
toolCall := map[string]interface{}{
"index": index,
"id": toolUseID,
"type": "function",
"function": map[string]interface{}{
"name": toolName,
"arguments": "",
},
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"tool_calls": []map[string]interface{}{toolCall},
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
return results
case "content_block_delta":
index := int(root.Get("index").Int())
deltaType := root.Get("delta.type").String()
if deltaType == "text_delta" {
// Text content delta
contentDelta := root.Get("delta.text").String()
if contentDelta != "" {
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"content": contentDelta,
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
} else if deltaType == "input_json_delta" {
// Tool input delta (streaming arguments)
partialJSON := root.Get("delta.partial_json").String()
if partialJSON != "" {
toolCall := map[string]interface{}{
"index": index,
"function": map[string]interface{}{
"arguments": partialJSON,
},
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"tool_calls": []map[string]interface{}{toolCall},
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
}
return results
case "content_block_stop":
// End of content block - no output needed for OpenAI format
return results
case "message_delta":
// Final message delta with stop_reason
stopReason := root.Get("delta.stop_reason").String()
if stopReason != "" {
finishReason := "stop"
if stopReason == "tool_use" {
finishReason = "tool_calls"
} else if stopReason == "end_turn" {
finishReason = "stop"
} else if stopReason == "max_tokens" {
finishReason = "length"
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{},
"finish_reason": finishReason,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
return results
case "message_stop":
// End of message - could emit [DONE] marker
return results
}
// Fallback: handle raw content for backward compatibility
var contentDelta string
if delta := root.Get("delta.text"); delta.Exists() {
contentDelta = delta.String()
} else if content := root.Get("content"); content.Exists() && root.Get("type").String() == "" {
contentDelta = content.String()
}
if contentDelta != "" {
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"content": contentDelta,
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
// Handle tool_use content blocks (Claude format) - fallback
toolUses := root.Get("delta.tool_use")
if !toolUses.Exists() {
toolUses = root.Get("tool_use")
}
if toolUses.Exists() && toolUses.IsObject() {
inputJSON := toolUses.Get("input").String()
if inputJSON == "" {
if inputObj := toolUses.Get("input"); inputObj.Exists() {
inputBytes, _ := json.Marshal(inputObj.Value())
inputJSON = string(inputBytes)
}
}
toolCall := map[string]interface{}{
"index": 0,
"id": toolUses.Get("id").String(),
"type": "function",
"function": map[string]interface{}{
"name": toolUses.Get("name").String(),
"arguments": inputJSON,
},
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion.chunk",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"delta": map[string]interface{}{
"tool_calls": []map[string]interface{}{toolCall},
},
"finish_reason": nil,
},
},
}
result, _ := json.Marshal(response)
results = append(results, string(result))
}
return results
}
// ConvertKiroResponseToOpenAINonStream converts Kiro non-streaming response to OpenAI format.
func ConvertKiroResponseToOpenAINonStream(ctx context.Context, model string, originalRequest, request, rawResponse []byte, param *any) string {
root := gjson.ParseBytes(rawResponse)
var content string
var toolCalls []map[string]interface{}
contentArray := root.Get("content")
if contentArray.IsArray() {
for _, item := range contentArray.Array() {
itemType := item.Get("type").String()
if itemType == "text" {
content += item.Get("text").String()
} else if itemType == "tool_use" {
// Convert Claude tool_use to OpenAI tool_calls format
inputJSON := item.Get("input").String()
if inputJSON == "" {
// If input is an object, marshal it
if inputObj := item.Get("input"); inputObj.Exists() {
inputBytes, _ := json.Marshal(inputObj.Value())
inputJSON = string(inputBytes)
}
}
toolCall := map[string]interface{}{
"id": item.Get("id").String(),
"type": "function",
"function": map[string]interface{}{
"name": item.Get("name").String(),
"arguments": inputJSON,
},
}
toolCalls = append(toolCalls, toolCall)
}
}
} else {
content = root.Get("content").String()
}
inputTokens := root.Get("usage.input_tokens").Int()
outputTokens := root.Get("usage.output_tokens").Int()
message := map[string]interface{}{
"role": "assistant",
"content": content,
}
// Add tool_calls if present
if len(toolCalls) > 0 {
message["tool_calls"] = toolCalls
}
finishReason := "stop"
if len(toolCalls) > 0 {
finishReason = "tool_calls"
}
response := map[string]interface{}{
"id": "chatcmpl-" + uuid.New().String()[:24],
"object": "chat.completion",
"created": time.Now().Unix(),
"model": model,
"choices": []map[string]interface{}{
{
"index": 0,
"message": message,
"finish_reason": finishReason,
},
},
"usage": map[string]interface{}{
"prompt_tokens": inputTokens,
"completion_tokens": outputTokens,
"total_tokens": inputTokens + outputTokens,
},
}
result, _ := json.Marshal(response)
return string(result)
}

View File

@@ -21,6 +21,7 @@ import (
"github.com/fsnotify/fsnotify"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"
"gopkg.in/yaml.v3"
@@ -176,6 +177,9 @@ func (w *Watcher) Start(ctx context.Context) error {
}
log.Debugf("watching auth directory: %s", w.authDir)
// Watch Kiro IDE token file directory for automatic token updates
w.watchKiroIDETokenFile()
// Start the event processing goroutine
go w.processEvents(ctx)
@@ -184,6 +188,31 @@ func (w *Watcher) Start(ctx context.Context) error {
return nil
}
// watchKiroIDETokenFile adds the Kiro IDE token file directory to the watcher.
// This enables automatic detection of token updates from Kiro IDE.
func (w *Watcher) watchKiroIDETokenFile() {
homeDir, err := os.UserHomeDir()
if err != nil {
log.Debugf("failed to get home directory for Kiro IDE token watch: %v", err)
return
}
// Kiro IDE stores tokens in ~/.aws/sso/cache/
kiroTokenDir := filepath.Join(homeDir, ".aws", "sso", "cache")
// Check if directory exists
if _, statErr := os.Stat(kiroTokenDir); os.IsNotExist(statErr) {
log.Debugf("Kiro IDE token directory does not exist: %s", kiroTokenDir)
return
}
if errAdd := w.watcher.Add(kiroTokenDir); errAdd != nil {
log.Debugf("failed to watch Kiro IDE token directory %s: %v", kiroTokenDir, errAdd)
return
}
log.Debugf("watching Kiro IDE token directory: %s", kiroTokenDir)
}
// Stop stops the file watcher
func (w *Watcher) Stop() error {
w.stopDispatch()
@@ -744,10 +773,20 @@ func (w *Watcher) handleEvent(event fsnotify.Event) {
isConfigEvent := event.Name == w.configPath && event.Op&configOps != 0
authOps := fsnotify.Create | fsnotify.Write | fsnotify.Remove | fsnotify.Rename
isAuthJSON := strings.HasPrefix(event.Name, w.authDir) && strings.HasSuffix(event.Name, ".json") && event.Op&authOps != 0
if !isConfigEvent && !isAuthJSON {
// Check for Kiro IDE token file changes
isKiroIDEToken := w.isKiroIDETokenFile(event.Name) && event.Op&authOps != 0
if !isConfigEvent && !isAuthJSON && !isKiroIDEToken {
// Ignore unrelated files (e.g., cookie snapshots *.cookie) and other noise.
return
}
// Handle Kiro IDE token file changes
if isKiroIDEToken {
w.handleKiroIDETokenChange(event)
return
}
now := time.Now()
log.Debugf("file system event detected: %s %s", event.Op.String(), event.Name)
@@ -805,6 +844,51 @@ func (w *Watcher) scheduleConfigReload() {
})
}
// isKiroIDETokenFile checks if the given path is the Kiro IDE token file.
func (w *Watcher) isKiroIDETokenFile(path string) bool {
// Check if it's the kiro-auth-token.json file in ~/.aws/sso/cache/
// Use filepath.ToSlash to ensure consistent separators across platforms (Windows uses backslashes)
normalized := filepath.ToSlash(path)
return strings.HasSuffix(normalized, "kiro-auth-token.json") && strings.Contains(normalized, ".aws/sso/cache")
}
// handleKiroIDETokenChange processes changes to the Kiro IDE token file.
// When the token file is updated by Kiro IDE, this triggers a reload of Kiro auth.
func (w *Watcher) handleKiroIDETokenChange(event fsnotify.Event) {
log.Debugf("Kiro IDE token file event detected: %s %s", event.Op.String(), event.Name)
if event.Op&(fsnotify.Remove|fsnotify.Rename) != 0 {
// Token file removed - wait briefly for potential atomic replace
time.Sleep(replaceCheckDelay)
if _, statErr := os.Stat(event.Name); statErr != nil {
log.Debugf("Kiro IDE token file removed: %s", event.Name)
return
}
}
// Try to load the updated token
tokenData, err := kiroauth.LoadKiroIDEToken()
if err != nil {
log.Debugf("failed to load Kiro IDE token after change: %v", err)
return
}
log.Infof("Kiro IDE token file updated, access token refreshed (provider: %s)", tokenData.Provider)
// Trigger auth state refresh to pick up the new token
w.refreshAuthState()
// Notify callback if set
w.clientsMutex.RLock()
cfg := w.config
w.clientsMutex.RUnlock()
if w.reloadCallback != nil && cfg != nil {
log.Debugf("triggering server update callback after Kiro IDE token change")
w.reloadCallback(cfg)
}
}
func (w *Watcher) reloadConfigIfChanged() {
data, err := os.ReadFile(w.configPath)
if err != nil {
@@ -1181,6 +1265,67 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
applyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
out = append(out, a)
}
// Kiro (AWS CodeWhisperer) -> synthesize auths
var kAuth *kiroauth.KiroAuth
if len(cfg.KiroKey) > 0 {
kAuth = kiroauth.NewKiroAuth(cfg)
}
for i := range cfg.KiroKey {
kk := cfg.KiroKey[i]
var accessToken, profileArn string
// Try to load from token file first
if kk.TokenFile != "" && kAuth != nil {
tokenData, err := kAuth.LoadTokenFromFile(kk.TokenFile)
if err != nil {
log.Warnf("failed to load kiro token file %s: %v", kk.TokenFile, err)
} else {
accessToken = tokenData.AccessToken
profileArn = tokenData.ProfileArn
}
}
// Override with direct config values if provided
if kk.AccessToken != "" {
accessToken = kk.AccessToken
}
if kk.ProfileArn != "" {
profileArn = kk.ProfileArn
}
if accessToken == "" {
log.Warnf("kiro config[%d] missing access_token, skipping", i)
continue
}
// profileArn is optional for AWS Builder ID users
id, token := idGen.next("kiro:token", accessToken, profileArn)
attrs := map[string]string{
"source": fmt.Sprintf("config:kiro[%s]", token),
"access_token": accessToken,
}
if profileArn != "" {
attrs["profile_arn"] = profileArn
}
if kk.Region != "" {
attrs["region"] = kk.Region
}
if kk.AgentTaskType != "" {
attrs["agent_task_type"] = kk.AgentTaskType
}
proxyURL := strings.TrimSpace(kk.ProxyURL)
a := &coreauth.Auth{
ID: id,
Provider: "kiro",
Label: "kiro-token",
Status: coreauth.StatusActive,
ProxyURL: proxyURL,
Attributes: attrs,
CreatedAt: now,
UpdatedAt: now,
}
out = append(out, a)
}
for i := range cfg.OpenAICompatibility {
compat := &cfg.OpenAICompatibility[i]
providerName := strings.ToLower(strings.TrimSpace(compat.Name))
@@ -1287,7 +1432,12 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
}
// Also synthesize auth entries directly from auth files (for OAuth/file-backed providers)
entries, _ := os.ReadDir(w.authDir)
log.Debugf("SnapshotCoreAuths: scanning auth directory: %s", w.authDir)
entries, readErr := os.ReadDir(w.authDir)
if readErr != nil {
log.Errorf("SnapshotCoreAuths: failed to read auth directory %s: %v", w.authDir, readErr)
}
log.Debugf("SnapshotCoreAuths: found %d entries in auth directory", len(entries))
for _, e := range entries {
if e.IsDir() {
continue
@@ -1306,9 +1456,20 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
continue
}
t, _ := metadata["type"].(string)
// Detect Kiro auth files by auth_method field (they don't have "type" field)
if t == "" {
if authMethod, _ := metadata["auth_method"].(string); authMethod == "builder-id" || authMethod == "social" {
t = "kiro"
log.Debugf("SnapshotCoreAuths: detected Kiro auth by auth_method: %s", name)
}
}
if t == "" {
log.Debugf("SnapshotCoreAuths: skipping file without type: %s", name)
continue
}
log.Debugf("SnapshotCoreAuths: processing auth file: %s (type=%s)", name, t)
provider := strings.ToLower(t)
if provider == "gemini" {
provider = "gemini-cli"
@@ -1317,6 +1478,12 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
if email, _ := metadata["email"].(string); email != "" {
label = email
}
// For Kiro, use provider field as label if available
if provider == "kiro" {
if kiroProvider, _ := metadata["provider"].(string); kiroProvider != "" {
label = fmt.Sprintf("kiro-%s", strings.ToLower(kiroProvider))
}
}
// Use relative path under authDir as ID to stay consistent with the file-based token store
id := full
if rel, errRel := filepath.Rel(w.authDir, full); errRel == nil && rel != "" {
@@ -1342,6 +1509,16 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
CreatedAt: now,
UpdatedAt: now,
}
// Set NextRefreshAfter for Kiro auth based on expires_at
if provider == "kiro" {
if expiresAtStr, ok := metadata["expires_at"].(string); ok && expiresAtStr != "" {
if expiresAt, parseErr := time.Parse(time.RFC3339, expiresAtStr); parseErr == nil {
// Refresh 30 minutes before expiry
a.NextRefreshAfter = expiresAt.Add(-30 * time.Minute)
}
}
}
applyAuthExcludedModelsMeta(a, cfg, nil, "oauth")
if provider == "gemini-cli" {
if virtuals := synthesizeGeminiVirtualAuths(a, metadata, now); len(virtuals) > 0 {

View File

@@ -8,6 +8,7 @@ import (
"fmt"
"net/http"
"strings"
"sync"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
@@ -50,7 +51,25 @@ type BaseAPIHandler struct {
Cfg *config.SDKConfig
// OpenAICompatProviders is a list of provider names for OpenAI compatibility.
OpenAICompatProviders []string
openAICompatProviders []string
openAICompatMutex sync.RWMutex
}
// GetOpenAICompatProviders safely returns a copy of the provider names
func (h *BaseAPIHandler) GetOpenAICompatProviders() []string {
h.openAICompatMutex.RLock()
defer h.openAICompatMutex.RUnlock()
result := make([]string, len(h.openAICompatProviders))
copy(result, h.openAICompatProviders)
return result
}
// SetOpenAICompatProviders safely sets the provider names
func (h *BaseAPIHandler) SetOpenAICompatProviders(providers []string) {
h.openAICompatMutex.Lock()
defer h.openAICompatMutex.Unlock()
h.openAICompatProviders = make([]string, len(providers))
copy(h.openAICompatProviders, providers)
}
// NewBaseAPIHandlers creates a new API handlers instance.
@@ -63,11 +82,12 @@ type BaseAPIHandler struct {
// Returns:
// - *BaseAPIHandler: A new API handlers instance
func NewBaseAPIHandlers(cfg *config.SDKConfig, authManager *coreauth.Manager, openAICompatProviders []string) *BaseAPIHandler {
return &BaseAPIHandler{
Cfg: cfg,
AuthManager: authManager,
OpenAICompatProviders: openAICompatProviders,
h := &BaseAPIHandler{
Cfg: cfg,
AuthManager: authManager,
}
h.SetOpenAICompatProviders(openAICompatProviders)
return h
}
// UpdateClients updates the handlers' client list and configuration.
@@ -363,7 +383,7 @@ func (h *BaseAPIHandler) parseDynamicModel(modelName string) (providerName, mode
}
// Check if the provider is a configured openai-compatibility provider
for _, pName := range h.OpenAICompatProviders {
for _, pName := range h.GetOpenAICompatProviders() {
if pName == providerPart {
return providerPart, modelPart, true
}

357
sdk/auth/kiro.go Normal file
View File

@@ -0,0 +1,357 @@
package auth
import (
"context"
"fmt"
"strings"
"time"
kiroauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/kiro"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
)
// extractKiroIdentifier extracts a meaningful identifier for file naming.
// Returns account name if provided, otherwise profile ARN ID.
// All extracted values are sanitized to prevent path injection attacks.
func extractKiroIdentifier(accountName, profileArn string) string {
// Priority 1: Use account name if provided
if accountName != "" {
return kiroauth.SanitizeEmailForFilename(accountName)
}
// Priority 2: Use profile ARN ID part (sanitized to prevent path injection)
if profileArn != "" {
parts := strings.Split(profileArn, "/")
if len(parts) >= 2 {
// Sanitize the ARN component to prevent path traversal
return kiroauth.SanitizeEmailForFilename(parts[len(parts)-1])
}
}
// Fallback: timestamp
return fmt.Sprintf("%d", time.Now().UnixNano()%100000)
}
// KiroAuthenticator implements OAuth authentication for Kiro with Google login.
type KiroAuthenticator struct{}
// NewKiroAuthenticator constructs a Kiro authenticator.
func NewKiroAuthenticator() *KiroAuthenticator {
return &KiroAuthenticator{}
}
// Provider returns the provider key for the authenticator.
func (a *KiroAuthenticator) Provider() string {
return "kiro"
}
// RefreshLead indicates how soon before expiry a refresh should be attempted.
func (a *KiroAuthenticator) RefreshLead() *time.Duration {
d := 30 * time.Minute
return &d
}
// Login performs OAuth login for Kiro with AWS Builder ID.
func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
if cfg == nil {
return nil, fmt.Errorf("kiro auth: configuration is required")
}
oauth := kiroauth.NewKiroOAuth(cfg)
// Use AWS Builder ID device code flow
tokenData, err := oauth.LoginWithBuilderID(ctx)
if err != nil {
return nil, fmt.Errorf("login failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
now := time.Now()
fileName := fmt.Sprintf("kiro-aws-%s.json", idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: "kiro-aws",
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"client_id": tokenData.ClientID,
"client_secret": tokenData.ClientSecret,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "aws-builder-id",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
if tokenData.Email != "" {
fmt.Printf("\n✓ Kiro authentication completed successfully! (Account: %s)\n", tokenData.Email)
} else {
fmt.Println("\n✓ Kiro authentication completed successfully!")
}
return record, nil
}
// LoginWithGoogle performs OAuth login for Kiro with Google.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (a *KiroAuthenticator) LoginWithGoogle(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
if cfg == nil {
return nil, fmt.Errorf("kiro auth: configuration is required")
}
oauth := kiroauth.NewKiroOAuth(cfg)
// Use Google OAuth flow with protocol handler
tokenData, err := oauth.LoginWithGoogle(ctx)
if err != nil {
return nil, fmt.Errorf("google login failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
now := time.Now()
fileName := fmt.Sprintf("kiro-google-%s.json", idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: "kiro-google",
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "google-oauth",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
if tokenData.Email != "" {
fmt.Printf("\n✓ Kiro Google authentication completed successfully! (Account: %s)\n", tokenData.Email)
} else {
fmt.Println("\n✓ Kiro Google authentication completed successfully!")
}
return record, nil
}
// LoginWithGitHub performs OAuth login for Kiro with GitHub.
// This uses a custom protocol handler (kiro://) to receive the callback.
func (a *KiroAuthenticator) LoginWithGitHub(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
if cfg == nil {
return nil, fmt.Errorf("kiro auth: configuration is required")
}
oauth := kiroauth.NewKiroOAuth(cfg)
// Use GitHub OAuth flow with protocol handler
tokenData, err := oauth.LoginWithGitHub(ctx)
if err != nil {
return nil, fmt.Errorf("github login failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
now := time.Now()
fileName := fmt.Sprintf("kiro-github-%s.json", idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: "kiro-github",
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "github-oauth",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
if tokenData.Email != "" {
fmt.Printf("\n✓ Kiro GitHub authentication completed successfully! (Account: %s)\n", tokenData.Email)
} else {
fmt.Println("\n✓ Kiro GitHub authentication completed successfully!")
}
return record, nil
}
// ImportFromKiroIDE imports token from Kiro IDE's token file.
func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.Config) (*coreauth.Auth, error) {
tokenData, err := kiroauth.LoadKiroIDEToken()
if err != nil {
return nil, fmt.Errorf("failed to load Kiro IDE token: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Extract email from JWT if not already set (for imported tokens)
if tokenData.Email == "" {
tokenData.Email = kiroauth.ExtractEmailFromJWT(tokenData.AccessToken)
}
// Extract identifier for file naming
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
// Sanitize provider to prevent path traversal (defense-in-depth)
provider := kiroauth.SanitizeEmailForFilename(strings.ToLower(strings.TrimSpace(tokenData.Provider)))
if provider == "" {
provider = "imported" // Fallback for legacy tokens without provider
}
now := time.Now()
fileName := fmt.Sprintf("kiro-%s-%s.json", provider, idPart)
record := &coreauth.Auth{
ID: fileName,
Provider: "kiro",
FileName: fileName,
Label: fmt.Sprintf("kiro-%s", provider),
Status: coreauth.StatusActive,
CreatedAt: now,
UpdatedAt: now,
Metadata: map[string]any{
"type": "kiro",
"access_token": tokenData.AccessToken,
"refresh_token": tokenData.RefreshToken,
"profile_arn": tokenData.ProfileArn,
"expires_at": tokenData.ExpiresAt,
"auth_method": tokenData.AuthMethod,
"provider": tokenData.Provider,
"email": tokenData.Email,
},
Attributes: map[string]string{
"profile_arn": tokenData.ProfileArn,
"source": "kiro-ide-import",
"email": tokenData.Email,
},
NextRefreshAfter: expiresAt.Add(-30 * time.Minute),
}
// Display the email if extracted
if tokenData.Email != "" {
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s, Account: %s)\n", tokenData.Provider, tokenData.Email)
} else {
fmt.Printf("\n✓ Imported Kiro token from IDE (Provider: %s)\n", tokenData.Provider)
}
return record, nil
}
// Refresh refreshes an expired Kiro token using AWS SSO OIDC.
func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, auth *coreauth.Auth) (*coreauth.Auth, error) {
if auth == nil || auth.Metadata == nil {
return nil, fmt.Errorf("invalid auth record")
}
refreshToken, ok := auth.Metadata["refresh_token"].(string)
if !ok || refreshToken == "" {
return nil, fmt.Errorf("refresh token not found")
}
clientID, _ := auth.Metadata["client_id"].(string)
clientSecret, _ := auth.Metadata["client_secret"].(string)
authMethod, _ := auth.Metadata["auth_method"].(string)
var tokenData *kiroauth.KiroTokenData
var err error
// Use SSO OIDC refresh for AWS Builder ID, otherwise use Kiro's OAuth refresh endpoint
if clientID != "" && clientSecret != "" && authMethod == "builder-id" {
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
} else {
// Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub)
oauth := kiroauth.NewKiroOAuth(cfg)
tokenData, err = oauth.RefreshToken(ctx, refreshToken)
}
if err != nil {
return nil, fmt.Errorf("token refresh failed: %w", err)
}
// Parse expires_at
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
if err != nil {
expiresAt = time.Now().Add(1 * time.Hour)
}
// Clone auth to avoid mutating the input parameter
updated := auth.Clone()
now := time.Now()
updated.UpdatedAt = now
updated.LastRefreshedAt = now
updated.Metadata["access_token"] = tokenData.AccessToken
updated.Metadata["refresh_token"] = tokenData.RefreshToken
updated.Metadata["expires_at"] = tokenData.ExpiresAt
updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization
updated.NextRefreshAfter = expiresAt.Add(-30 * time.Minute)
return updated, nil
}

View File

@@ -74,3 +74,16 @@ func (m *Manager) Login(ctx context.Context, provider string, cfg *config.Config
}
return record, savedPath, nil
}
// SaveAuth persists an auth record directly without going through the login flow.
func (m *Manager) SaveAuth(record *coreauth.Auth, cfg *config.Config) (string, error) {
if m.store == nil {
return "", fmt.Errorf("no store configured")
}
if cfg != nil {
if dirSetter, ok := m.store.(interface{ SetBaseDir(string) }); ok {
dirSetter.SetBaseDir(cfg.AuthDir)
}
}
return m.store.Save(context.Background(), record)
}

View File

@@ -14,6 +14,7 @@ func init() {
registerRefreshLead("gemini", func() Authenticator { return NewGeminiAuthenticator() })
registerRefreshLead("gemini-cli", func() Authenticator { return NewGeminiAuthenticator() })
registerRefreshLead("antigravity", func() Authenticator { return NewAntigravityAuthenticator() })
registerRefreshLead("kiro", func() Authenticator { return NewKiroAuthenticator() })
}
func registerRefreshLead(provider string, factory func() Authenticator) {

View File

@@ -379,6 +379,8 @@ func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) {
s.coreManager.RegisterExecutor(executor.NewQwenExecutor(s.cfg))
case "iflow":
s.coreManager.RegisterExecutor(executor.NewIFlowExecutor(s.cfg))
case "kiro":
s.coreManager.RegisterExecutor(executor.NewKiroExecutor(s.cfg))
default:
providerKey := strings.ToLower(strings.TrimSpace(a.Provider))
if providerKey == "" {
@@ -721,6 +723,9 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
case "iflow":
models = registry.GetIFlowModels()
models = applyExcludedModels(models, excluded)
case "kiro":
models = registry.GetKiroModels()
models = applyExcludedModels(models, excluded)
default:
// Handle OpenAI-compatibility providers by name using config
if s.cfg != nil {