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.
This commit is contained in:
yuechenglong.5
2026-01-24 20:02:09 +08:00
parent 9fccc86b71
commit 8f780e7280
2 changed files with 45 additions and 36 deletions

View File

@@ -6,6 +6,7 @@ import (
"time" "time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
@@ -49,6 +50,14 @@ func (m *RefreshManager) Initialize(baseDir string, cfg *config.Config) error {
return nil 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 存储库 // 创建 token 存储库
repo := NewFileTokenRepository(baseDir) repo := NewFileTokenRepository(baseDir)

View File

@@ -791,28 +791,28 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth.
_ = httpResp.Body.Close() _ = httpResp.Body.Close()
appendAPIResponseChunk(ctx, e.cfg, respBody) appendAPIResponseChunk(ctx, e.cfg, respBody)
if attempt < maxRetries { log.Warnf("kiro: received 401 error, attempting token refresh")
log.Warnf("kiro: received 401 error, attempting token refresh and retry (attempt %d/%d)", attempt+1, maxRetries+1) 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 refreshedAuth != nil {
if refreshErr != nil { auth = refreshedAuth
log.Errorf("kiro: token refresh failed: %v", refreshErr) // Persist the refreshed auth to file so subsequent requests use it
return resp, statusErr{code: httpResp.StatusCode, msg: string(respBody)} 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)
if refreshedAuth != nil { // Rebuild payload with new profile ARN if changed
auth = refreshedAuth kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
// Persist the refreshed auth to file so subsequent requests use it if attempt < maxRetries {
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { log.Infof("kiro: token refreshed successfully, retrying request (attempt %d/%d)", attempt+1, maxRetries+1)
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")
continue 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)) 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() _ = httpResp.Body.Close()
appendAPIResponseChunk(ctx, e.cfg, respBody) appendAPIResponseChunk(ctx, e.cfg, respBody)
if attempt < maxRetries { log.Warnf("kiro: stream received 401 error, attempting token refresh")
log.Warnf("kiro: stream received 401 error, attempting token refresh and retry (attempt %d/%d)", attempt+1, maxRetries+1) 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 refreshedAuth != nil {
if refreshErr != nil { auth = refreshedAuth
log.Errorf("kiro: token refresh failed: %v", refreshErr) // Persist the refreshed auth to file so subsequent requests use it
return nil, statusErr{code: httpResp.StatusCode, msg: string(respBody)} 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)
if refreshedAuth != nil { // Rebuild payload with new profile ARN if changed
auth = refreshedAuth kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
// Persist the refreshed auth to file so subsequent requests use it if attempt < maxRetries {
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil { log.Infof("kiro: token refreshed successfully, retrying stream request (attempt %d/%d)", attempt+1, maxRetries+1)
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")
continue continue
} }
log.Infof("kiro: token refreshed successfully, no retries remaining")
} }
log.Warnf("kiro stream error, status: 401, body: %s", string(respBody)) log.Warnf("kiro stream error, status: 401, body: %s", string(respBody))