ai-gateway-go/internal/provider/runtime/resolver.go

package runtime

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"hash/fnv"
	"log/slog"
	"strings"
	"sync/atomic"
	"time"

	"aigateway.local/core/internal/gateway"
	"aigateway.local/core/internal/provider"
	provideropenai "aigateway.local/core/internal/provider/openai"
	"github.com/redis/go-redis/v9"
)

const changeChannel = "gateway:providers:changed"

type snapshot struct {
	adapters    map[string]gateway.ResolvedAdapter
	defaultCode string
	routes      map[string][]compiledRoute
	knownModels map[string]bool
}

type routeConditions struct {
	Endpoints []string `json:"endpoints"`
	APIKeyIDs []string `json:"api_key_ids"`
	TenantIDs []string `json:"tenant_ids"`
}

type compiledRoute struct {
	providerCode string
	targetModel  string
	weight       int
	priority     int
	conditions   routeConditions
}

type Resolver struct {
	repository *provider.Repository
	cipher     *provider.CredentialCipher
	fallback   gateway.ResolvedAdapter
	interval   time.Duration
	logger     *slog.Logger
	redis      *redis.Client
	snapshot   atomic.Pointer[snapshot]
}

func (r *Resolver) SetNotificationClient(client *redis.Client) {
	r.redis = client
}

func (r *Resolver) Notify(ctx context.Context) error {
	if r.redis == nil {
		return nil
	}
	return r.redis.Publish(ctx, changeChannel, time.Now().UTC().Format(time.RFC3339Nano)).Err()
}

func NewResolver(repository *provider.Repository, cipher *provider.CredentialCipher, fallback provider.Adapter, interval time.Duration, logger *slog.Logger) *Resolver {
	capabilities := make(map[provider.Capability]bool)
	if fallback != nil {
		for _, capability := range fallback.Capabilities() {
			capabilities[capability] = true
		}
	}
	resolver := &Resolver{
		repository: repository, cipher: cipher, interval: interval, logger: logger,
		fallback: gateway.ResolvedAdapter{Code: "environment", Adapter: fallback, Capabilities: capabilities},
	}
	initial := &snapshot{adapters: make(map[string]gateway.ResolvedAdapter), routes: make(map[string][]compiledRoute), knownModels: make(map[string]bool)}
	if fallback != nil {
		initial.adapters["environment"] = resolver.fallback
		initial.defaultCode = "environment"
	}
	resolver.snapshot.Store(initial)
	return resolver
}

func (r *Resolver) ResolveModelRoute(query gateway.ModelRouteQuery) (gateway.ModelRouteResult, error) {
	current := r.snapshot.Load()
	if current == nil {
		return gateway.ModelRouteResult{}, gateway.ErrProviderUnavailable
	}
	candidates := current.routes[query.Model]
	known := current.knownModels[query.Model]
	matched := make([]compiledRoute, 0, len(candidates))
	highestPriority := -100001
	for _, route := range candidates {
		if query.ProviderCode != "" && route.providerCode != query.ProviderCode || !route.conditions.matches(query) {
			continue
		}
		if route.priority > highestPriority {
			highestPriority = route.priority
			matched = matched[:0]
		}
		if route.priority == highestPriority {
			matched = append(matched, route)
		}
	}
	if len(matched) == 0 {
		return gateway.ModelRouteResult{Known: known}, nil
	}
	total := 0
	for _, route := range matched {
		total += route.weight
	}
	hasher := fnv.New64a()
	_, _ = hasher.Write([]byte(query.Seed + "\x00" + query.Model))
	selected := int(hasher.Sum64() % uint64(total))
	chosen := matched[len(matched)-1]
	for _, route := range matched {
		if selected < route.weight {
			chosen = route
			break
		}
		selected -= route.weight
	}
	resolved, ok := current.adapters[chosen.providerCode]
	if !ok {
		return gateway.ModelRouteResult{}, gateway.ErrProviderUnavailable
	}
	return gateway.ModelRouteResult{ResolvedAdapter: resolved, TargetModel: chosen.targetModel, Matched: true, Known: true}, nil
}

func (r *Resolver) ModelRoutingEnabled() bool {
	current := r.snapshot.Load()
	return current != nil && len(current.knownModels) > 0
}

func (c routeConditions) matches(query gateway.ModelRouteQuery) bool {
	return containsOrEmpty(c.Endpoints, query.Endpoint) && containsOrEmpty(c.APIKeyIDs, query.APIKeyID) && containsOrEmpty(c.TenantIDs, query.TenantID)
}

func containsOrEmpty(values []string, target string) bool {
	if len(values) == 0 {
		return true
	}
	for _, value := range values {
		if value == target {
			return true
		}
	}
	return false
}

func (r *Resolver) Resolve(code string) (gateway.ResolvedAdapter, error) {
	current := r.snapshot.Load()
	if current == nil {
		return gateway.ResolvedAdapter{}, gateway.ErrProviderUnavailable
	}
	code = strings.ToLower(strings.TrimSpace(code))
	if code == "" {
		if current.defaultCode == "" {
			return gateway.ResolvedAdapter{}, gateway.ErrProviderUnavailable
		}
		code = current.defaultCode
	}
	resolved, ok := current.adapters[code]
	if !ok {
		return gateway.ResolvedAdapter{}, gateway.ErrProviderNotFound
	}
	return resolved, nil
}

func (r *Resolver) Reload(ctx context.Context) error {
	records, err := r.repository.List(ctx)
	if err != nil {
		return err
	}
	routeRecords, err := r.repository.ListModelRoutes(ctx)
	if err != nil {
		return err
	}
	next := &snapshot{adapters: make(map[string]gateway.ResolvedAdapter), routes: make(map[string][]compiledRoute), knownModels: make(map[string]bool)}
	previous := r.snapshot.Load()
	enabledCount := 0
	explicitDefault := false
	for _, record := range records {
		if !record.Enabled {
			continue
		}
		enabledCount++
		resolved, defaultProvider, err := r.build(record)
		if err != nil {
			r.logger.Error("provider snapshot entry rejected", "provider", record.Code, "revision", record.Revision, "error", err)
			if previous != nil {
				if lastValid, ok := previous.adapters[record.Code]; ok {
					next.adapters[record.Code] = lastValid
					if next.defaultCode == "" || previous.defaultCode == record.Code {
						next.defaultCode = record.Code
					}
				}
			}
			continue
		}
		next.adapters[record.Code] = resolved
		if defaultProvider {
			if explicitDefault {
				return errors.New("multiple enabled providers declare config.default=true; retaining previous snapshot")
			}
			explicitDefault = true
			next.defaultCode = record.Code
		} else if next.defaultCode == "" {
			next.defaultCode = record.Code
		}
	}
	if enabledCount > 0 && len(next.adapters) == 0 {
		return errors.New("all enabled provider records are invalid; retaining previous snapshot")
	}
	if len(next.adapters) == 0 {
		if r.fallback.Adapter != nil {
			next.adapters["environment"] = r.fallback
			next.defaultCode = "environment"
		}
	}
	for _, route := range routeRecords {
		if !route.Enabled {
			continue
		}
		if _, ok := next.adapters[route.ProviderCode]; !ok {
			continue
		}
		// Only routes that can actually be served make the model "known".
		// Marking it before the Enabled check above meant a disabled route
		// still advertised the model as known, and the proxy hard-blocks any
		// known-but-unmatched model with a 400 — so disabling the only route
		// for a model made that model unusable on every other provider.
		next.knownModels[route.SourceModel] = true
		var conditions routeConditions
		if err := json.Unmarshal(route.Conditions, &conditions); err != nil {
			r.logger.Error("model route snapshot entry rejected", "route", route.ID, "error", err)
			continue
		}
		next.routes[route.SourceModel] = append(next.routes[route.SourceModel], compiledRoute{
			providerCode: route.ProviderCode, targetModel: route.TargetModel, weight: route.Weight,
			priority: route.Priority, conditions: conditions,
		})
	}
	r.snapshot.Store(next)
	r.logger.Info("provider snapshot refreshed", "providers", len(next.adapters), "model_routes", len(routeRecords), "default", next.defaultCode)
	return nil
}

func (r *Resolver) Run(ctx context.Context) {
	if err := r.Reload(ctx); err != nil {
		r.logger.Warn("initial provider snapshot refresh failed; using last valid snapshot", "error", err)
	}
	interval := r.interval
	if interval <= 0 {
		interval = 5 * time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	var changes <-chan *redis.Message
	var subscription *redis.PubSub
	if r.redis != nil {
		subscription = r.redis.Subscribe(ctx, changeChannel)
		defer subscription.Close()
		changes = subscription.Channel(redis.WithChannelSize(32))
	}
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			refreshCtx, cancel := context.WithTimeout(ctx, interval)
			err := r.Reload(refreshCtx)
			cancel()
			if err != nil && ctx.Err() == nil {
				r.logger.Warn("provider snapshot refresh failed; retaining last valid snapshot", "error", err)
			}
		case _, ok := <-changes:
			if !ok {
				changes = nil
				continue
			}
			refreshCtx, cancel := context.WithTimeout(ctx, interval)
			err := r.Reload(refreshCtx)
			cancel()
			if err != nil && ctx.Err() == nil {
				r.logger.Warn("provider notification refresh failed; retaining last valid snapshot", "error", err)
			}
		}
	}
}

func (r *Resolver) build(record provider.Record) (gateway.ResolvedAdapter, bool, error) {
	plaintext, err := r.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
	if err != nil {
		return gateway.ResolvedAdapter{}, false, fmt.Errorf("decrypt credentials: %w", err)
	}
	var credentials provider.Credentials
	if err := json.Unmarshal(plaintext, &credentials); err != nil {
		return gateway.ResolvedAdapter{}, false, fmt.Errorf("decode credentials: %w", err)
	}
	var adapter provider.Adapter
	switch record.Adapter {
	case "openai-compatible":
		adapter, err = provideropenai.New(record.BaseURL, credentials.APIKey)
	default:
		err = fmt.Errorf("unsupported adapter %q", record.Adapter)
	}
	if err != nil {
		return gateway.ResolvedAdapter{}, false, err
	}
	capabilities := make(map[provider.Capability]bool, len(record.Capabilities))
	for _, value := range record.Capabilities {
		capability := provider.Capability(value)
		switch capability {
		case provider.CapabilityChat, provider.CapabilityResponses, provider.CapabilityEmbeddings, provider.CapabilityMessages, provider.CapabilityModels:
			capabilities[capability] = true
		default:
			return gateway.ResolvedAdapter{}, false, fmt.Errorf("unsupported capability %q", value)
		}
	}
	var configuration struct {
		Default bool `json:"default"`
	}
	if len(record.Config) > 0 {
		if err := json.Unmarshal(record.Config, &configuration); err != nil {
			return gateway.ResolvedAdapter{}, false, fmt.Errorf("decode config: %w", err)
		}
	}
	return gateway.ResolvedAdapter{Code: record.Code, Revision: record.Revision, Adapter: adapter, Capabilities: capabilities}, configuration.Default, nil
}