CLIProxyAPIPlus/internal/auth/kiro/refresh_utils.go

// Package kiro provides refresh utilities for Kiro token management.
package kiro

import (
	"context"
	"fmt"
	"time"

	log "github.com/sirupsen/logrus"
)

// RefreshResult contains the result of a token refresh attempt.
type RefreshResult struct {
	TokenData    *KiroTokenData
	Error        error
	UsedFallback bool // True if we used the existing token as fallback
}

// RefreshWithGracefulDegradation attempts to refresh a token with graceful degradation.
// If refresh fails but the existing access token is still valid, it returns the existing token.
// This matches kiro-openai-gateway's behavior for better reliability.
//
// Parameters:
//   - ctx: Context for the request
//   - refreshFunc: Function to perform the actual refresh
//   - existingAccessToken: Current access token (for fallback)
//   - expiresAt: Expiration time of the existing token
//
// Returns:
//   - RefreshResult containing the new or existing token data
func RefreshWithGracefulDegradation(
	ctx context.Context,
	refreshFunc func(ctx context.Context) (*KiroTokenData, error),
	existingAccessToken string,
	expiresAt time.Time,
) RefreshResult {
	// Try to refresh the token
	newTokenData, err := refreshFunc(ctx)
	if err == nil {
		return RefreshResult{
			TokenData:    newTokenData,
			Error:        nil,
			UsedFallback: false,
		}
	}

	// Refresh failed - check if we can use the existing token
	log.Warnf("kiro: token refresh failed: %v", err)

	// Check if existing token is still valid (not expired)
	if existingAccessToken != "" && time.Now().Before(expiresAt) {
		remainingTime := time.Until(expiresAt)
		log.Warnf("kiro: using existing access token (expires in %v). Will retry refresh later.", remainingTime.Round(time.Second))

		return RefreshResult{
			TokenData: &KiroTokenData{
				AccessToken: existingAccessToken,
				ExpiresAt:   expiresAt.Format(time.RFC3339),
			},
			Error:        nil,
			UsedFallback: true,
		}
	}

	// Token is expired and refresh failed - return the error
	return RefreshResult{
		TokenData:    nil,
		Error:        fmt.Errorf("token refresh failed and existing token is expired: %w", err),
		UsedFallback: false,
	}
}

// IsTokenExpiringSoon checks if a token is expiring within the given threshold.
// Default threshold is 5 minutes if not specified.
func IsTokenExpiringSoon(expiresAt time.Time, threshold time.Duration) bool {
	if threshold == 0 {
		threshold = 5 * time.Minute
	}
	return time.Now().Add(threshold).After(expiresAt)
}

// IsTokenExpired checks if a token has already expired.
func IsTokenExpired(expiresAt time.Time) bool {
	return time.Now().After(expiresAt)
}

// ParseExpiresAt parses an expiration time string in RFC3339 format.
// Returns zero time if parsing fails.
func ParseExpiresAt(expiresAtStr string) time.Time {
	if expiresAtStr == "" {
		return time.Time{}
	}
	t, err := time.Parse(time.RFC3339, expiresAtStr)
	if err != nil {
		log.Debugf("kiro: failed to parse expiresAt '%s': %v", expiresAtStr, err)
		return time.Time{}
	}
	return t
}

// RefreshConfig contains configuration for token refresh behavior.
type RefreshConfig struct {
	// MaxRetries is the maximum number of refresh attempts (default: 1)
	MaxRetries int
	// RetryDelay is the delay between retry attempts (default: 1 second)
	RetryDelay time.Duration
	// RefreshThreshold is how early to refresh before expiration (default: 5 minutes)
	RefreshThreshold time.Duration
	// EnableGracefulDegradation allows using existing token if refresh fails (default: true)
	EnableGracefulDegradation bool
}

// DefaultRefreshConfig returns the default refresh configuration.
func DefaultRefreshConfig() RefreshConfig {
	return RefreshConfig{
		MaxRetries:                1,
		RetryDelay:                time.Second,
		RefreshThreshold:          5 * time.Minute,
		EnableGracefulDegradation: true,
	}
}

// RefreshWithRetry attempts to refresh a token with retry logic.
func RefreshWithRetry(
	ctx context.Context,
	refreshFunc func(ctx context.Context) (*KiroTokenData, error),
	config RefreshConfig,
) (*KiroTokenData, error) {
	var lastErr error

	maxAttempts := config.MaxRetries + 1
	if maxAttempts < 1 {
		maxAttempts = 1
	}

	for attempt := 1; attempt <= maxAttempts; attempt++ {
		tokenData, err := refreshFunc(ctx)
		if err == nil {
			if attempt > 1 {
				log.Infof("kiro: token refresh succeeded on attempt %d", attempt)
			}
			return tokenData, nil
		}

		lastErr = err
		log.Warnf("kiro: token refresh attempt %d/%d failed: %v", attempt, maxAttempts, err)

		// Don't sleep after the last attempt
		if attempt < maxAttempts {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(config.RetryDelay):
			}
		}
	}

	return nil, fmt.Errorf("token refresh failed after %d attempts: %w", maxAttempts, lastErr)
}