From 8f780e7280469a67b74e94af305a5f5ee10d497e Mon Sep 17 00:00:00 2001 From: "yuechenglong.5" Date: Sat, 24 Jan 2026 20:02:09 +0800 Subject: [PATCH] fix(kiro): always attempt token refresh on 401 before checking retry count Refactor 401 error handling in both executeWithRetry and executeStreamWithRetry to always attempt token refresh regardless of remaining retry attempts. Previously, token refresh was only attempted when retries remained, which could leave valid refreshed tokens unused. Also add auth directory resolution in RefreshManager.Initialize to properly resolve the base directory path before creating the token repository. --- internal/auth/kiro/refresh_manager.go | 9 +++ internal/runtime/executor/kiro_executor.go | 72 +++++++++++----------- 2 files changed, 45 insertions(+), 36 deletions(-) diff --git a/internal/auth/kiro/refresh_manager.go b/internal/auth/kiro/refresh_manager.go index 05e27a54..5330c5e1 100644 --- a/internal/auth/kiro/refresh_manager.go +++ b/internal/auth/kiro/refresh_manager.go @@ -6,6 +6,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + "github.com/router-for-me/CLIProxyAPI/v6/internal/util" log "github.com/sirupsen/logrus" ) @@ -49,6 +50,14 @@ func (m *RefreshManager) Initialize(baseDir string, cfg *config.Config) error { return nil } + resolvedBaseDir, err := util.ResolveAuthDir(baseDir) + if err != nil { + log.Warnf("refresh manager: failed to resolve auth directory %s: %v", baseDir, err) + } + if resolvedBaseDir != "" { + baseDir = resolvedBaseDir + } + // 创建 token 存储库 repo := NewFileTokenRepository(baseDir) diff --git a/internal/runtime/executor/kiro_executor.go b/internal/runtime/executor/kiro_executor.go index ed6014a2..57574268 100644 --- a/internal/runtime/executor/kiro_executor.go +++ b/internal/runtime/executor/kiro_executor.go @@ -791,28 +791,28 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth. _ = httpResp.Body.Close() appendAPIResponseChunk(ctx, e.cfg, respBody) - if attempt < maxRetries { - log.Warnf("kiro: received 401 error, attempting token refresh and retry (attempt %d/%d)", attempt+1, maxRetries+1) + log.Warnf("kiro: received 401 error, attempting token refresh") + refreshedAuth, refreshErr := e.Refresh(ctx, auth) + if refreshErr != nil { + log.Errorf("kiro: token refresh failed: %v", refreshErr) + return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)} + } - refreshedAuth, refreshErr := e.Refresh(ctx, auth) - if refreshErr != nil { - log.Errorf("kiro: token refresh failed: %v", refreshErr) - return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)} + if refreshedAuth != nil { + auth = refreshedAuth + // Persist the refreshed auth to file so subsequent requests use it + if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { + log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) + // Continue anyway - the token is valid for this request } - - if refreshedAuth != nil { - auth = refreshedAuth - // Persist the refreshed auth to file so subsequent requests use it - if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { - log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) - // Continue anyway - the token is valid for this request - } - accessToken, profileArn = kiroCredentials(auth) - // Rebuild payload with new profile ARN if changed - kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers) - log.Infof("kiro: token refreshed successfully, retrying request") + accessToken, profileArn = kiroCredentials(auth) + // Rebuild payload with new profile ARN if changed + kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers) + if attempt < maxRetries { + log.Infof("kiro: token refreshed successfully, retrying request (attempt %d/%d)", attempt+1, maxRetries+1) continue } + log.Infof("kiro: token refreshed successfully, no retries remaining") } log.Warnf("kiro request error, status: 401, body: %s", summarizeErrorBody(httpResp.Header.Get("Content-Type"), respBody)) @@ -1199,28 +1199,28 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox _ = httpResp.Body.Close() appendAPIResponseChunk(ctx, e.cfg, respBody) - if attempt < maxRetries { - log.Warnf("kiro: stream received 401 error, attempting token refresh and retry (attempt %d/%d)", attempt+1, maxRetries+1) + log.Warnf("kiro: stream received 401 error, attempting token refresh") + refreshedAuth, refreshErr := e.Refresh(ctx, auth) + if refreshErr != nil { + log.Errorf("kiro: token refresh failed: %v", refreshErr) + return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)} + } - refreshedAuth, refreshErr := e.Refresh(ctx, auth) - if refreshErr != nil { - log.Errorf("kiro: token refresh failed: %v", refreshErr) - return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)} + if refreshedAuth != nil { + auth = refreshedAuth + // Persist the refreshed auth to file so subsequent requests use it + if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { + log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) + // Continue anyway - the token is valid for this request } - - if refreshedAuth != nil { - auth = refreshedAuth - // Persist the refreshed auth to file so subsequent requests use it - if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { - log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr) - // Continue anyway - the token is valid for this request - } - accessToken, profileArn = kiroCredentials(auth) - // Rebuild payload with new profile ARN if changed - kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers) - log.Infof("kiro: token refreshed successfully, retrying stream request") + accessToken, profileArn = kiroCredentials(auth) + // Rebuild payload with new profile ARN if changed + kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers) + if attempt < maxRetries { + log.Infof("kiro: token refreshed successfully, retrying stream request (attempt %d/%d)", attempt+1, maxRetries+1) continue } + log.Infof("kiro: token refreshed successfully, no retries remaining") } log.Warnf("kiro stream error, status: 401, body: %s", string(respBody))