package channel import ( "context" "fmt" "log" "math/rand" "opencatd-open/internal/dao" "opencatd-open/internal/store" "opencatd-open/internal/pkg/crypto" "sync" "time" ) type Service struct { channelDAO *dao.ChannelDAO modelDAO *dao.ModelDAO // Health tracking mu sync.RWMutex healthStatus map[uint64]*channelHealth } type channelHealth struct { status string consecutive int lastCheck time.Time cooldown time.Time } func NewService(channelDAO *dao.ChannelDAO, modelDAO *dao.ModelDAO) *Service { return &Service{ channelDAO: channelDAO, modelDAO: modelDAO, healthStatus: make(map[uint64]*channelHealth), } } // SelectChannel selects the best channel for a given model using weighted random selection func (s *Service) SelectChannel(ctx context.Context, modelName string) (*store.Channel, error) { channels, err := s.channelDAO.GetEnabledChannelsByModel(modelName) if err != nil { return nil, fmt.Errorf("failed to get channels for model %s: %w", modelName, err) } if len(channels) == 0 { return nil, fmt.Errorf("no enabled channels for model: %s", modelName) } // Filter out unhealthy channels candidates := s.filterHealthy(channels) if len(candidates) == 0 { // If all channels are unhealthy, try the first one anyway candidates = channels[:1] } // Weighted random selection totalWeight := 0 for _, ch := range candidates { totalWeight += ch.Weight } if totalWeight == 0 { return candidates[0], nil } r := rand.Intn(totalWeight) for _, ch := range candidates { r -= ch.Weight if r < 0 { return ch, nil } } return candidates[0], nil } // GetChannelByKeyID decrypts the API key for a channel func (s *Service) GetChannelByKeyID(ctx context.Context, channelID uint64) (*store.Channel, error) { ch, err := s.channelDAO.GetByID(channelID) if err != nil { return nil, err } return ch, nil } // GetAPIKey decrypts the channel's API key func (s *Service) GetAPIKey(ch *store.Channel) (string, error) { return crypto.Decrypt(ch.APIKeyEnc) } // RecordSuccess records a successful request to a channel func (s *Service) RecordSuccess(channelID uint64) { s.mu.Lock() defer s.mu.Unlock() h := s.getOrCreateHealth(channelID) h.consecutive = 0 h.status = store.ChannelHealthHealthy h.lastCheck = time.Now() } // RecordFailure records a failed request to a channel func (s *Service) RecordFailure(channelID uint64) { s.mu.Lock() defer s.mu.Unlock() h := s.getOrCreateHealth(channelID) h.consecutive++ h.lastCheck = time.Now() if h.consecutive >= 3 { h.status = store.ChannelHealthDegraded h.cooldown = time.Now().Add(5 * time.Minute) } if h.consecutive >= 5 { h.status = store.ChannelHealthCooldown h.cooldown = time.Now().Add(15 * time.Minute) } } // RecordTimeout records a timeout to a channel func (s *Service) RecordTimeout(channelID uint64) { s.RecordFailure(channelID) } func (s *Service) getOrCreateHealth(channelID uint64) *channelHealth { h, ok := s.healthStatus[channelID] if !ok { h = &channelHealth{ status: store.ChannelHealthHealthy, } s.healthStatus[channelID] = h } return h } func (s *Service) filterHealthy(channels []*store.Channel) []*store.Channel { s.mu.RLock() defer s.mu.RUnlock() var healthy []*store.Channel now := time.Now() for _, ch := range channels { h, ok := s.healthStatus[ch.ID] if !ok { healthy = append(healthy, ch) continue } // Check if cooldown has expired if now.After(h.cooldown) && h.cooldown.IsZero() == false { h.consecutive = 0 h.status = store.ChannelHealthHealthy healthy = append(healthy, ch) continue } if h.status == store.ChannelHealthHealthy || h.status == store.ChannelHealthDegraded { healthy = append(healthy, ch) } } return healthy } // GetHealthStatus returns the health status of a channel func (s *Service) GetHealthStatus(channelID uint64) string { s.mu.RLock() defer s.mu.RUnlock() h, ok := s.healthStatus[channelID] if !ok { return store.ChannelHealthHealthy } return h.status } // ChannelCandidate represents a channel with its resolved API key type ChannelCandidate struct { Channel *store.Channel APIKey string Format string } // SelectCandidates returns candidates for a model, sorted by priority func (s *Service) SelectCandidates(ctx context.Context, modelName string, preferredFormat string) ([]ChannelCandidate, error) { channels, err := s.channelDAO.GetEnabledChannelsByModel(modelName) if err != nil { return nil, err } var candidates []ChannelCandidate for _, ch := range channels { // Check if channel supports the preferred format formats := ch.FormatsEffective() supported := false for _, f := range formats { if f == preferredFormat || preferredFormat == "" { supported = true break } } if !supported { continue } apiKey, err := crypto.Decrypt(ch.APIKeyEnc) if err != nil { log.Printf("Failed to decrypt API key for channel %s: %v", ch.Name, err) continue } candidates = append(candidates, ChannelCandidate{ Channel: ch, APIKey: apiKey, Format: preferredFormat, }) } return candidates, nil }