openteam/server/internal/channel/health.go

package channel

import (
	"context"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/openteam/server/internal/pkg/crypto"
	"github.com/openteam/server/internal/store"
	"gorm.io/gorm"
)

// HealthConfig 健康检查参数。
type HealthConfig struct {
	Interval      time.Duration // 探测周期
	FailThreshold int           // 连续失败 N 次进入 cooldown
	Timeout       time.Duration // 单次探测超时
}

// HealthMonitor 定时探测渠道健康状态。
type HealthMonitor struct {
	db   *gorm.DB
	enc  *crypto.Encryptor
	cfg  HealthConfig
	hc   *http.Client
	mu   sync.Mutex
	fail map[uint64]int // 渠道连续失败次数
	// OnChange 健康状态变化后的回调（目录缓存据此立即失效，不必等 TTL）。
	OnChange func()
}

func NewHealthMonitor(db *gorm.DB, enc *crypto.Encryptor, cfg HealthConfig) *HealthMonitor {
	if cfg.Interval <= 0 {
		cfg.Interval = 60 * time.Second
	}
	if cfg.FailThreshold <= 0 {
		cfg.FailThreshold = 2
	}
	if cfg.Timeout <= 0 {
		cfg.Timeout = 5 * time.Second
	}
	return &HealthMonitor{
		db: db, enc: enc, cfg: cfg,
		hc:   &http.Client{Timeout: cfg.Timeout},
		fail: map[uint64]int{},
	}
}

// Start 启动后台探测循环（ctx 取消即停止）。
func (h *HealthMonitor) Start(ctx context.Context) {
	go func() {
		ticker := time.NewTicker(h.cfg.Interval)
		defer ticker.Stop()
		h.probeAll()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				h.probeAll()
			}
		}
	}()
	log.Printf("health: monitor started (interval=%s, threshold=%d)", h.cfg.Interval, h.cfg.FailThreshold)
}

// probeAll 探测所有启用渠道并更新健康状态。
func (h *HealthMonitor) probeAll() {
	var chs []store.Channel
	if err := h.db.Where("enabled = ?", true).Find(&chs).Error; err != nil {
		return
	}
	for i := range chs {
		h.probe(&chs[i])
	}
}

// probe 单渠道探测：GET {base}/v1/models。
func (h *HealthMonitor) probe(ch *store.Channel) {
	key, err := h.enc.Decrypt(ch.APIKeyEnc)
	if err != nil {
		h.record(ch, false)
		return
	}
	url := ch.UpstreamURL("", "/models")
	req, _ := http.NewRequest(http.MethodGet, url, nil)
	req.Header.Set("Authorization", "Bearer "+key)
	req.Header.Set("Accept", "application/json")
	// Anthropic 官方 API 只认 x-api-key（Bearer 被忽略），探测必须带上才不会误判不健康。
	// 两个头同时发：官方 API 忽略 Bearer，OpenAI 兼容渠道忽略 x-api-key，互不干扰。
	for _, f := range ch.FormatsEffective() {
		if f == store.FormatMessages {
			req.Header.Set("x-api-key", key)
			req.Header.Set("anthropic-version", "2023-06-01")
			break
		}
	}
	resp, err := h.hc.Do(req)
	ok := err == nil && resp.StatusCode >= 200 && resp.StatusCode < 300
	if resp != nil {
		resp.Body.Close()
	}
	h.record(ch, ok)
}

// record 记录一次探测结果，按阈值切换健康状态。
func (h *HealthMonitor) record(ch *store.Channel, ok bool) {
	h.mu.Lock()
	defer h.mu.Unlock()

	prev := ch.HealthStatus
	if ok {
		h.fail[ch.ID] = 0
		if prev != store.ChannelHealthHealthy {
			h.setStatus(ch, store.ChannelHealthHealthy)
			log.Printf("health: channel %q recovered -> healthy", ch.Name)
		}
		return
	}
	h.fail[ch.ID]++
	if h.fail[ch.ID] >= h.cfg.FailThreshold && prev != store.ChannelHealthCooldown {
		h.setStatus(ch, store.ChannelHealthCooldown)
		log.Printf("health: channel %q -> cooldown (%d consecutive failures)", ch.Name, h.fail[ch.ID])
	}
}

func (h *HealthMonitor) setStatus(ch *store.Channel, status string) {
	h.db.Model(&store.Channel{}).Where("id = ?", ch.ID).Update("health_status", status)
	ch.HealthStatus = status
	if h.OnChange != nil {
		h.OnChange()
	}
}