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 { // 防御:weight<=0 的行(绕过管理端校验直接入库的脏数据)不得参与 // 权重池,否则 total=0 时取模除零 panic,同一请求 ID 将永远 500。 if route.weight <= 0 { continue } total += route.weight } weighted := make([]compiledRoute, 0, len(matched)) for _, route := range matched { if route.weight > 0 { weighted = append(weighted, route) } } if len(weighted) == 0 { return gateway.ModelRouteResult{Known: known}, nil } hasher := fnv.New64a() _, _ = hasher.Write([]byte(query.Seed + "\x00" + query.Model)) selected := int(hasher.Sum64() % uint64(total)) chosen := weighted[len(weighted)-1] for _, route := range weighted { 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 }