github-actions[bot]
Update from GitHub Actions
6fefda3
package jetbrains
import (
"context"
"fmt"
"github.com/go-resty/resty/v2"
"jetbrains-ai-proxy/internal/balancer"
"jetbrains-ai-proxy/internal/config"
"jetbrains-ai-proxy/internal/types"
"jetbrains-ai-proxy/internal/utils"
"log"
"sync"
)
var (
jwtBalancer balancer.JWTBalancer
healthChecker *balancer.HealthChecker
initOnce sync.Once
configManager *config.Manager
)
// InitializeFromConfig 从配置管理器初始化JWT负载均衡器
func InitializeFromConfig() error {
var initErr error
initOnce.Do(func() {
configManager = config.GetGlobalConfig()
// 加载配置
if err := configManager.LoadConfig(); err != nil {
initErr = fmt.Errorf("failed to load config: %v", err)
return
}
// 获取配置
cfg := configManager.GetConfig()
tokens := configManager.GetJWTTokens()
if len(tokens) == 0 {
initErr = fmt.Errorf("no JWT tokens configured")
return
}
// 创建负载均衡器
jwtBalancer = balancer.NewJWTBalancer(tokens, cfg.LoadBalanceStrategy)
// 创建并启动健康检查器
healthChecker = balancer.NewHealthChecker(jwtBalancer)
if cfg.HealthCheckInterval > 0 {
healthChecker.SetCheckInterval(cfg.HealthCheckInterval)
}
healthChecker.Start()
log.Printf("JWT balancer initialized from config:")
log.Printf(" - Tokens: %d", len(tokens))
log.Printf(" - Strategy: %s", cfg.LoadBalanceStrategy)
log.Printf(" - Health check interval: %v", cfg.HealthCheckInterval)
})
return initErr
}
// InitializeBalancer 初始化JWT负载均衡器(向后兼容)
func InitializeBalancer(tokens []string, strategy string) error {
if len(tokens) == 0 {
return fmt.Errorf("no JWT tokens provided")
}
var balanceStrategy config.LoadBalanceStrategy
switch strategy {
case "random":
balanceStrategy = config.Random
case "round_robin", "":
balanceStrategy = config.RoundRobin
default:
balanceStrategy = config.RoundRobin
}
// 创建负载均衡器
jwtBalancer = balancer.NewJWTBalancer(tokens, balanceStrategy)
// 创建并启动健康检查器
healthChecker = balancer.NewHealthChecker(jwtBalancer)
healthChecker.Start()
log.Printf("JWT balancer initialized with %d tokens, strategy: %s", len(tokens), string(balanceStrategy))
return nil
}
// ReloadConfig 重新加载配置
func ReloadConfig() error {
if configManager == nil {
return fmt.Errorf("config manager not initialized")
}
// 重新加载配置
if err := configManager.LoadConfig(); err != nil {
return fmt.Errorf("failed to reload config: %v", err)
}
// 获取新配置
cfg := configManager.GetConfig()
tokens := configManager.GetJWTTokens()
if len(tokens) == 0 {
return fmt.Errorf("no JWT tokens in reloaded config")
}
// 更新负载均衡器
if jwtBalancer != nil {
jwtBalancer.RefreshTokens(tokens)
}
// 更新健康检查间隔
if healthChecker != nil && cfg.HealthCheckInterval > 0 {
healthChecker.SetCheckInterval(cfg.HealthCheckInterval)
}
log.Printf("Config reloaded successfully:")
log.Printf(" - Tokens: %d", len(tokens))
log.Printf(" - Strategy: %s", cfg.LoadBalanceStrategy)
return nil
}
// StopBalancer 停止负载均衡器
func StopBalancer() {
if healthChecker != nil {
healthChecker.Stop()
}
}
// GetConfigManager 获取配置管理器
func GetConfigManager() *config.Manager {
return configManager
}
func SendJetbrainsRequest(ctx context.Context, req *types.JetbrainsRequest) (*resty.Response, error) {
// 获取一个可用的JWT token
token, err := jwtBalancer.GetToken()
if err != nil {
log.Printf("failed to get JWT token: %v", err)
return nil, fmt.Errorf("no available JWT tokens: %v", err)
}
resp, err := utils.RestySSEClient.R().
SetContext(ctx).
SetHeader(types.JwtTokenKey, token).
SetDoNotParseResponse(true).
SetBody(req).
Post(types.ChatStreamV7)
if err != nil {
log.Printf("jetbrains ai req error: %v", err)
// 标记token为不健康
jwtBalancer.MarkTokenUnhealthy(token)
return nil, err
}
// 检查响应状态码
if resp.StatusCode() == 401 {
// 401表示token无效,标记为不健康
jwtBalancer.MarkTokenUnhealthy(token)
log.Printf("JWT token invalid (401): %s...", token[:min(len(token), 10)])
return nil, fmt.Errorf("JWT token invalid")
} else if resp.StatusCode() == 200 {
// 成功响应,确保token标记为健康
jwtBalancer.MarkTokenHealthy(token)
}
return resp, nil
}
// GetBalancerStats 获取负载均衡器统计信息
func GetBalancerStats() (int, int) {
if jwtBalancer == nil {
return 0, 0
}
return jwtBalancer.GetHealthyTokenCount(), jwtBalancer.GetTotalTokenCount()
}
// min 辅助函数
func min(a, b int) int {
if a < b {
return a
}
return b
}