From a9ee971e1c30499c409319471b7dea8b23f0c2ee Mon Sep 17 00:00:00 2001 From: "781456868@qq.com" Date: Tue, 20 Jan 2026 21:57:45 +0800 Subject: [PATCH] fix(kiro): improve auto-refresh and IDC auth file handling Amp-Thread-ID: https://ampcode.com/threads/T-019bdb94-80e3-7302-be0f-a69937826d13 Co-authored-by: Amp --- internal/auth/kiro/aws_auth.go | 20 +++++++++++++++++ internal/auth/kiro/oauth_web.go | 25 +++++++++++----------- internal/runtime/executor/kiro_executor.go | 14 ++++++------ sdk/auth/filestore.go | 4 ++-- sdk/auth/kiro.go | 20 ++++++++--------- sdk/cliproxy/auth/conductor.go | 2 +- 6 files changed, 53 insertions(+), 32 deletions(-) diff --git a/internal/auth/kiro/aws_auth.go b/internal/auth/kiro/aws_auth.go index 53c77a8b..d082f274 100644 --- a/internal/auth/kiro/aws_auth.go +++ b/internal/auth/kiro/aws_auth.go @@ -280,6 +280,11 @@ func (k *KiroAuth) CreateTokenStorage(tokenData *KiroTokenData) *KiroTokenStorag AuthMethod: tokenData.AuthMethod, Provider: tokenData.Provider, LastRefresh: time.Now().Format(time.RFC3339), + ClientID: tokenData.ClientID, + ClientSecret: tokenData.ClientSecret, + Region: tokenData.Region, + StartURL: tokenData.StartURL, + Email: tokenData.Email, } } @@ -311,4 +316,19 @@ func (k *KiroAuth) UpdateTokenStorage(storage *KiroTokenStorage, tokenData *Kiro storage.AuthMethod = tokenData.AuthMethod storage.Provider = tokenData.Provider storage.LastRefresh = time.Now().Format(time.RFC3339) + if tokenData.ClientID != "" { + storage.ClientID = tokenData.ClientID + } + if tokenData.ClientSecret != "" { + storage.ClientSecret = tokenData.ClientSecret + } + if tokenData.Region != "" { + storage.Region = tokenData.Region + } + if tokenData.StartURL != "" { + storage.StartURL = tokenData.StartURL + } + if tokenData.Email != "" { + storage.Email = tokenData.Email + } } diff --git a/internal/auth/kiro/oauth_web.go b/internal/auth/kiro/oauth_web.go index 81c24393..6e4269c5 100644 --- a/internal/auth/kiro/oauth_web.go +++ b/internal/auth/kiro/oauth_web.go @@ -377,17 +377,18 @@ func (h *OAuthWebHandler) pollForToken(ctx context.Context, session *webAuthSess email := FetchUserEmailWithFallback(ctx, h.cfg, tokenResp.AccessToken) tokenData := &KiroTokenData{ - AccessToken: tokenResp.AccessToken, - RefreshToken: tokenResp.RefreshToken, - ProfileArn: profileArn, - ExpiresAt: expiresAt.Format(time.RFC3339), - AuthMethod: session.authMethod, - Provider: "AWS", - ClientID: session.clientID, - ClientSecret: session.clientSecret, - Email: email, - Region: session.region, - } + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + ProfileArn: profileArn, + ExpiresAt: expiresAt.Format(time.RFC3339), + AuthMethod: session.authMethod, + Provider: "AWS", + ClientID: session.clientID, + ClientSecret: session.clientSecret, + Email: email, + Region: session.region, + StartURL: session.startURL, + } h.mu.Lock() session.status = statusSuccess @@ -828,7 +829,7 @@ func (h *OAuthWebHandler) handleImportToken(c *gin.Context) { // handleManualRefresh handles manual token refresh requests from the web UI. // This allows users to trigger a token refresh when needed, without waiting -// for the automatic 5-second check and 10-minute-before-expiry refresh cycle. +// for the automatic 30-second check and 20-minute-before-expiry refresh cycle. // Uses the same refresh logic as kiro_executor.Refresh for consistency. func (h *OAuthWebHandler) handleManualRefresh(c *gin.Context) { authDir := "" diff --git a/internal/runtime/executor/kiro_executor.go b/internal/runtime/executor/kiro_executor.go index b0c14c61..cab4bcd6 100644 --- a/internal/runtime/executor/kiro_executor.go +++ b/internal/runtime/executor/kiro_executor.go @@ -3513,14 +3513,14 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c // Also check if expires_at is now in the future with sufficient buffer if expiresAt, ok := auth.Metadata["expires_at"].(string); ok { if expTime, err := time.Parse(time.RFC3339, expiresAt); err == nil { - // If token expires more than 5 minutes from now, it's still valid - if time.Until(expTime) > 5*time.Minute { + // If token expires more than 20 minutes from now, it's still valid + if time.Until(expTime) > 20*time.Minute { log.Debugf("kiro executor: token is still valid (expires in %v), skipping refresh", time.Until(expTime)) // CRITICAL FIX: Set NextRefreshAfter to prevent frequent refresh checks - // Without this, shouldRefresh() will return true again in 5 seconds + // Without this, shouldRefresh() will return true again in 30 seconds updated := auth.Clone() - // Set next refresh to 5 minutes before expiry, or at least 30 seconds from now - nextRefresh := expTime.Add(-5 * time.Minute) + // Set next refresh to 20 minutes before expiry, or at least 30 seconds from now + nextRefresh := expTime.Add(-20 * time.Minute) minNextRefresh := time.Now().Add(30 * time.Second) if nextRefresh.Before(minNextRefresh) { nextRefresh = minNextRefresh @@ -3626,9 +3626,9 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c updated.Attributes["profile_arn"] = tokenData.ProfileArn } - // NextRefreshAfter is aligned with RefreshLead (5min) + // NextRefreshAfter is aligned with RefreshLead (20min) if expiresAt, parseErr := time.Parse(time.RFC3339, tokenData.ExpiresAt); parseErr == nil { - updated.NextRefreshAfter = expiresAt.Add(-5 * time.Minute) + updated.NextRefreshAfter = expiresAt.Add(-20 * time.Minute) } log.Infof("kiro executor: token refreshed successfully, expires at %s", tokenData.ExpiresAt) diff --git a/sdk/auth/filestore.go b/sdk/auth/filestore.go index 0010be7d..9a288b10 100644 --- a/sdk/auth/filestore.go +++ b/sdk/auth/filestore.go @@ -217,11 +217,11 @@ func (s *FileTokenStore) readAuthFile(path, baseDir string) (*cliproxyauth.Auth, } id := s.idFor(path, baseDir) - // Calculate NextRefreshAfter from expires_at (10 minutes before expiry) + // Calculate NextRefreshAfter from expires_at (20 minutes before expiry) var nextRefreshAfter time.Time if expiresAtStr, ok := metadata["expires_at"].(string); ok && expiresAtStr != "" { if expiresAt, err := time.Parse(time.RFC3339, expiresAtStr); err == nil { - nextRefreshAfter = expiresAt.Add(-10 * time.Minute) + nextRefreshAfter = expiresAt.Add(-20 * time.Minute) } } diff --git a/sdk/auth/kiro.go b/sdk/auth/kiro.go index b0687eba..f66be461 100644 --- a/sdk/auth/kiro.go +++ b/sdk/auth/kiro.go @@ -52,9 +52,9 @@ func (a *KiroAuthenticator) Provider() string { } // RefreshLead indicates how soon before expiry a refresh should be attempted. -// Set to 10 minutes for proactive refresh before token expiry. +// Set to 20 minutes for proactive refresh before token expiry. func (a *KiroAuthenticator) RefreshLead() *time.Duration { - d := 10 * time.Minute + d := 20 * time.Minute return &d } @@ -132,8 +132,8 @@ func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, UpdatedAt: now, Metadata: metadata, Attributes: attributes, - // NextRefreshAfter: 10 minutes before expiry - NextRefreshAfter: expiresAt.Add(-10 * time.Minute), + // NextRefreshAfter: 20 minutes before expiry + NextRefreshAfter: expiresAt.Add(-20 * time.Minute), } if tokenData.Email != "" { @@ -214,8 +214,8 @@ func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.C "source": "aws-builder-id-authcode", "email": tokenData.Email, }, - // NextRefreshAfter: 10 minutes before expiry - NextRefreshAfter: expiresAt.Add(-10 * time.Minute), + // NextRefreshAfter: 20 minutes before expiry + NextRefreshAfter: expiresAt.Add(-20 * time.Minute), } if tokenData.Email != "" { @@ -298,8 +298,8 @@ func (a *KiroAuthenticator) ImportFromKiroIDE(ctx context.Context, cfg *config.C "email": tokenData.Email, "region": tokenData.Region, }, - // NextRefreshAfter: 10 minutes before expiry - NextRefreshAfter: expiresAt.Add(-10 * time.Minute), + // NextRefreshAfter: 20 minutes before expiry + NextRefreshAfter: expiresAt.Add(-20 * time.Minute), } // Display the email if extracted @@ -367,8 +367,8 @@ func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, aut updated.Metadata["refresh_token"] = tokenData.RefreshToken updated.Metadata["expires_at"] = tokenData.ExpiresAt updated.Metadata["last_refresh"] = now.Format(time.RFC3339) // For double-check optimization - // NextRefreshAfter: 10 minutes before expiry - updated.NextRefreshAfter = expiresAt.Add(-10 * time.Minute) + // NextRefreshAfter: 20 minutes before expiry + updated.NextRefreshAfter = expiresAt.Add(-20 * time.Minute) return updated, nil } diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 83769198..5f553bdd 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -47,7 +47,7 @@ type RefreshEvaluator interface { } const ( - refreshCheckInterval = 5 * time.Second + refreshCheckInterval = 30 * time.Second refreshPendingBackoff = time.Minute refreshFailureBackoff = 1 * time.Minute quotaBackoffBase = time.Second