Luminary/internal/gateway/failover.go

package gateway

import (
	"context"
	"fmt"
	"sort"
	"time"

	"github.com/rose_cat707/luminary/internal/balancer"
	"github.com/rose_cat707/luminary/internal/limiter"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/provider"
	"github.com/rose_cat707/luminary/internal/store/cache"
)

type attempt struct {
	provider    *model.Provider
	key         *model.ProviderKey
	apiKey      string
}

func (s *Service) buildAttempts(providers []model.Provider, modelName string, route *routeTarget) ([]attempt, error) {
	var attempts []attempt
	var decryptFailures int
	var eligibleKeys int
	for i := range providers {
		p := providers[i]
		keys := s.listEligibleKeys(p.ID, modelName, route)
		eligibleKeys += len(keys)
		for _, k := range keys {
			plain, err := s.DecryptAPIKey(k.APIKeyEnc)
			if err != nil {
				decryptFailures++
				continue
			}
			attempts = append(attempts, attempt{provider: &p, key: &k, apiKey: plain})
		}
	}
	if len(attempts) == 0 {
		if decryptFailures > 0 {
			return nil, fmt.Errorf("no available provider keys: %d key(s) failed to decrypt (check encryption_key in system settings)", decryptFailures)
		}
		if eligibleKeys == 0 {
			return nil, fmt.Errorf("no available provider keys: none match model %q or are disabled/rate-limited", modelName)
		}
		return nil, fmt.Errorf("no available provider keys")
	}
	sort.SliceStable(attempts, func(i, j int) bool {
		return attempts[i].key.Weight > attempts[j].key.Weight
	})
	return attempts, nil
}

func (s *Service) listEligibleKeys(providerID uint, modelName string, route *routeTarget) []model.ProviderKey {
	if route != nil && route.ProviderKeyID > 0 && route.ProviderID == providerID {
		if k, ok := s.cache.GetProviderKey(route.ProviderKeyID); ok && k.Enabled && limiter.IsProviderKeyAvailable(k) {
			return []model.ProviderKey{k}
		}
	}
	var candidates []balancer.KeyCandidate
	for _, k := range s.cache.ListProviderKeys(providerID) {
		if !limiter.IsProviderKeyAvailable(k) {
			continue
		}
		allowed := cache.ParseJSONList(k.ModelsJSON)
		if modelName != "" && !cache.MatchesList(allowed, "*") {
			upstream := s.resolveUpstreamModelID(providerID, modelName)
			if !cache.MatchesList(allowed, modelName) && !cache.MatchesList(allowed, upstream) {
				continue
			}
		}
		candidates = append(candidates, balancer.KeyCandidate{Key: k, Weight: k.Weight})
	}
	// collect unique keys by weighted random sampling order
	seen := map[uint]bool{}
	var out []model.ProviderKey
	for len(out) < len(candidates) {
		selected := balancer.SelectWeighted(candidates)
		if selected == nil {
			break
		}
		if seen[selected.ID] {
			// remove from candidates
			filtered := candidates[:0]
			for _, c := range candidates {
				if c.Key.ID != selected.ID {
					filtered = append(filtered, c)
				}
			}
			candidates = filtered
			if len(candidates) == 0 {
				break
			}
			continue
		}
		seen[selected.ID] = true
		out = append(out, *selected)
		filtered := candidates[:0]
		for _, c := range candidates {
			if c.Key.ID != selected.ID {
				filtered = append(filtered, c)
			}
		}
		candidates = filtered
	}
	return out
}

func (s *Service) forwardWithFailover(ctx context.Context, attempts []attempt, endpoint provider.EndpointType, body []byte) (*provider.ChatResponse, *attempt, error) {
	var lastErr error
	var lastAtt *attempt
	var lastResp *provider.ChatResponse
	maxRetries := s.settings.GatewayMaxRetries()
	if maxRetries < 1 {
		maxRetries = 1
	}
	backoffMs := s.settings.GatewayRetryBackoffMs()
	for _, att := range attempts {
		if !s.canServeEndpoint(att.provider, endpoint) {
			continue
		}
		for try := 0; try < maxRetries; try++ {
			if try > 0 {
				time.Sleep(time.Duration(backoffMs) * time.Millisecond)
			}
			resp, err := s.forward(ctx, att.provider, att.apiKey, endpoint, body)
			cur := att
			lastAtt = &cur
			if resp != nil {
				lastResp = resp
			}
			if err != nil {
				lastErr = err
				continue
			}
			if resp == nil {
				lastErr = fmt.Errorf("empty upstream response")
				continue
			}
			if resp.StatusCode < 400 {
				return resp, &att, nil
			}
			lastErr = fmt.Errorf("%s", describeForwardError(nil, resp))
			if passThroughUpstreamStatus(resp.StatusCode) {
				return resp, &att, nil
			}
			if tryNextKeyOnUpstreamStatus(resp.StatusCode) {
				break // try next provider key
			}
			if retryableUpstreamStatus(resp.StatusCode) {
				continue // retry same key
			}
			return resp, &att, nil
		}
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("all providers failed")
	}
	return lastResp, lastAtt, lastErr
}