CLIProxyAPI/sdk/cliproxy/auth/selector_rr.go

package auth

import (
	"context"
	"strings"

	conversation "github.com/router-for-me/CLIProxyAPI/v6/internal/provider/gemini-web/conversation"
	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
	log "github.com/sirupsen/logrus"
)

const (
	geminiWebProviderKey = "gemini-web"
)

type geminiWebStickySelector struct {
	base Selector
}

func NewGeminiWebStickySelector(base Selector) Selector {
	if selector, ok := base.(*geminiWebStickySelector); ok {
		return selector
	}
	if base == nil {
		base = &RoundRobinSelector{}
	}
	return &geminiWebStickySelector{base: base}
}

func (m *Manager) EnableGeminiWebStickySelector() {
	if m == nil {
		return
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.selector.(*geminiWebStickySelector); ok {
		return
	}
	m.selector = NewGeminiWebStickySelector(m.selector)
}

func (s *geminiWebStickySelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
	if !strings.EqualFold(provider, geminiWebProviderKey) {
		if opts.Metadata != nil {
			delete(opts.Metadata, conversation.MetadataMatchKey)
		}
		return s.base.Pick(ctx, provider, model, opts, auths)
	}

	messages := extractGeminiWebMessages(opts.Metadata)
	if len(messages) >= 2 {
		normalizedModel := conversation.NormalizeModel(model)
		candidates := conversation.BuildLookupHashes(normalizedModel, messages)
		for _, candidate := range candidates {
			record, ok, err := conversation.LookupMatch(candidate.Hash)
			if err != nil {
				log.Warnf("gemini-web selector: lookup failed for hash %s: %v", candidate.Hash, err)
				continue
			}
			if !ok {
				continue
			}
			label := strings.TrimSpace(record.AccountLabel)
			if label == "" {
				continue
			}
			auth := findAuthByLabel(auths, label)
			if auth != nil {
				if opts.Metadata != nil {
					opts.Metadata[conversation.MetadataMatchKey] = &conversation.MatchResult{
						Hash:   candidate.Hash,
						Record: record,
						Model:  normalizedModel,
					}
				}
				return auth, nil
			}
			_ = conversation.RemoveMatchForLabel(candidate.Hash, label)
		}
	}

	return s.base.Pick(ctx, provider, model, opts, auths)
}

func extractGeminiWebMessages(metadata map[string]any) []conversation.Message {
	if metadata == nil {
		return nil
	}
	raw, ok := metadata[conversation.MetadataMessagesKey]
	if !ok {
		return nil
	}
	switch v := raw.(type) {
	case []conversation.Message:
		return v
	case *[]conversation.Message:
		if v == nil {
			return nil
		}
		return *v
	default:
		return nil
	}
}

func findAuthByLabel(auths []*Auth, label string) *Auth {
	if len(auths) == 0 {
		return nil
	}
	normalized := strings.ToLower(strings.TrimSpace(label))
	for _, auth := range auths {
		if auth == nil {
			continue
		}
		if strings.ToLower(strings.TrimSpace(auth.Label)) == normalized {
			return auth
		}
		if auth.Metadata != nil {
			if v, ok := auth.Metadata["label"].(string); ok && strings.ToLower(strings.TrimSpace(v)) == normalized {
				return auth
			}
		}
	}
	return nil
}