sub2api/backend/internal/service/kimi_token_provider.go

package service

import (
	"context"
	"errors"
	"log/slog"
	"strconv"
	"strings"
	"time"
)

const (
	kimiTokenRefreshSkew = 3 * time.Minute
	kimiTokenCacheSkew   = 5 * time.Minute
)

// KimiTokenProvider manages access_token for Kimi OAuth accounts.
type KimiTokenProvider struct {
	accountRepo     AccountRepository
	tokenCache      GeminiTokenCache
	kimiOAuthService *KimiOAuthService
	refreshAPI      *OAuthRefreshAPI
	executor        OAuthRefreshExecutor
	refreshPolicy   ProviderRefreshPolicy
}

// NewKimiTokenProvider creates a new KimiTokenProvider.
func NewKimiTokenProvider(
	accountRepo AccountRepository,
	tokenCache GeminiTokenCache,
	kimiOAuthService *KimiOAuthService,
) *KimiTokenProvider {
	return &KimiTokenProvider{
		accountRepo:      accountRepo,
		tokenCache:       tokenCache,
		kimiOAuthService: kimiOAuthService,
		refreshPolicy:    KimiProviderRefreshPolicy(),
	}
}

// SetRefreshAPI injects unified OAuth refresh API and executor.
func (p *KimiTokenProvider) SetRefreshAPI(api *OAuthRefreshAPI, executor OAuthRefreshExecutor) {
	p.refreshAPI = api
	p.executor = executor
}

// SetRefreshPolicy injects caller-side refresh policy.
func (p *KimiTokenProvider) SetRefreshPolicy(policy ProviderRefreshPolicy) {
	p.refreshPolicy = policy
}

// GetAccessToken returns a valid access token for the given Kimi account.
func (p *KimiTokenProvider) GetAccessToken(ctx context.Context, account *Account) (string, error) {
	if account == nil {
		return "", errors.New("account is nil")
	}
	if account.Platform != PlatformKimi || account.Type != AccountTypeOAuth {
		return "", errors.New("not a kimi oauth account")
	}

	cacheKey := KimiTokenCacheKey(account)

	// 1) Try cache first.
	if p.tokenCache != nil {
		if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" {
			return token, nil
		}
	}

	// 2) Refresh if needed (pre-expiry skew).
	expiresAt := account.GetCredentialAsTime("expires_at")
	needsRefresh := expiresAt == nil || time.Until(*expiresAt) <= kimiTokenRefreshSkew

	if needsRefresh && p.refreshAPI != nil && p.executor != nil {
		result, err := p.refreshAPI.RefreshIfNeeded(ctx, account, p.executor, kimiTokenRefreshSkew)
		if err != nil {
			if p.refreshPolicy.OnRefreshError == ProviderRefreshErrorReturn {
				return "", err
			}
		} else if result.LockHeld {
			if p.refreshPolicy.OnLockHeld == ProviderLockHeldWaitForCache && p.tokenCache != nil {
				if token, cacheErr := p.tokenCache.GetAccessToken(ctx, cacheKey); cacheErr == nil && strings.TrimSpace(token) != "" {
					return token, nil
				}
			}
			slog.Debug("kimi_token_lock_held_use_old", "account_id", account.ID)
		} else {
			account = result.Account
			expiresAt = account.GetCredentialAsTime("expires_at")
		}
	} else if needsRefresh && p.tokenCache != nil {
		locked, lockErr := p.tokenCache.AcquireRefreshLock(ctx, cacheKey, 30*time.Second)
		if lockErr == nil && locked {
			defer func() { _ = p.tokenCache.ReleaseRefreshLock(ctx, cacheKey) }()
		} else if lockErr != nil {
			slog.Warn("kimi_token_lock_failed", "account_id", account.ID, "error", lockErr)
		}
	}

	accessToken := account.GetCredential("access_token")
	if strings.TrimSpace(accessToken) == "" {
		return "", errors.New("access_token not found in credentials")
	}

	// 3) Populate cache with TTL.
	if p.tokenCache != nil {
		latestAccount, isStale := CheckTokenVersion(ctx, account, p.accountRepo)
		if isStale && latestAccount != nil {
			slog.Debug("kimi_token_version_stale_use_latest", "account_id", account.ID)
			accessToken = latestAccount.GetCredential("access_token")
			if strings.TrimSpace(accessToken) == "" {
				return "", errors.New("access_token not found after version check")
			}
		} else {
			ttl := 30 * time.Minute
			if expiresAt != nil {
				until := time.Until(*expiresAt)
				switch {
				case until > kimiTokenCacheSkew:
					ttl = until - kimiTokenCacheSkew
				case until > 0:
					ttl = until
				default:
					ttl = time.Minute
				}
			}
			_ = p.tokenCache.SetAccessToken(ctx, cacheKey, accessToken, ttl)
		}
	}

	return accessToken, nil
}

// KimiTokenCacheKey returns the cache key for a Kimi account.
func KimiTokenCacheKey(account *Account) string {
	return "kimi:account:" + strconv.FormatInt(account.ID, 10)
}

// KimiProviderRefreshPolicy returns the default refresh policy for Kimi.
func KimiProviderRefreshPolicy() ProviderRefreshPolicy {
	return ProviderRefreshPolicy{
		OnRefreshError: ProviderRefreshErrorReturn,
		OnLockHeld:     ProviderLockHeldWaitForCache,
	}
}

// KimiTokenRefresher implements OAuthRefreshExecutor for Kimi.
type KimiTokenRefresher struct {
	kimiOAuthService *KimiOAuthService
}

// NewKimiTokenRefresher creates a new KimiTokenRefresher.
func NewKimiTokenRefresher(kimiOAuthService *KimiOAuthService) *KimiTokenRefresher {
	return &KimiTokenRefresher{kimiOAuthService: kimiOAuthService}
}

// CanRefresh checks if this refresher can handle the given account.
func (r *KimiTokenRefresher) CanRefresh(account *Account) bool {
	return account != nil && account.Platform == PlatformKimi && account.Type == AccountTypeOAuth
}

// NeedsRefresh checks if the account needs token refresh.
func (r *KimiTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.Duration) bool {
	if account == nil {
		return false
	}
	expiresAt := account.GetCredentialAsTime("expires_at")
	return expiresAt == nil || time.Until(*expiresAt) <= refreshWindow
}

// Refresh performs the token refresh for a Kimi account.
func (r *KimiTokenRefresher) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
	if account == nil || r.kimiOAuthService == nil {
		return nil, errors.New("kimi token refresher not initialized")
	}

	refreshToken := account.GetCredential("refresh_token")
	if strings.TrimSpace(refreshToken) == "" {
		return nil, errors.New("no refresh token available")
	}

	var proxyURL string
	if account.ProxyID != nil && account.Proxy != nil {
		proxyURL = account.Proxy.URL()
	}

	tokenInfo, err := r.kimiOAuthService.RefreshToken(ctx, refreshToken, proxyURL)
	if err != nil {
		return nil, err
	}

	return r.kimiOAuthService.BuildAccountCredentials(tokenInfo), nil
}

// CacheKey returns the cache key for distributed locking.
func (r *KimiTokenRefresher) CacheKey(account *Account) string {
	return KimiTokenCacheKey(account)
}