internal/registry/model_registry.go
Rama main · réplica de lectura del repositorio original.
// Package registry provides centralized model management for all AI service providers.
// It implements a dynamic model registry with reference counting to track active clients
// and automatically hide models when no clients are available or when quota is exceeded.
package registry
import (
"context"
"encoding/json"
"fmt"
"sort"
"strings"
"sync"
"sync/atomic"
"time"
misc "github.com/router-for-me/CLIProxyAPI/v8/internal/misc"
log "github.com/sirupsen/logrus"
)
// OpenAIImageModelType marks models that are callable through OpenAI-compatible image endpoints.
const OpenAIImageModelType = "openai-image"
const (
DefaultClaudeMaxInputTokens = 200000
DefaultClaudeMaxOutputTokens = 64000
)
// NativeCapabilities contains tri-state native capability metadata from the catalog.
type NativeCapabilities struct {
// WebSearch reports explicit per-model support for native web search.
// nil means the catalog has not established support either way.
WebSearch *bool `json:"web_search,omitempty"`
// ToolSearch reports whether every credential behind a Home route can take
// Codex tool_search as Anthropic native tool search. Home computes it per
// route; nil means unknown and is treated as unsupported.
ToolSearch *bool `json:"tool_search,omitempty"`
}
// ModelInfo represents information about an available model
type ModelInfo struct {
// ID is the unique identifier for the model
ID string `json:"id"`
// MetadataModelID identifies the canonical model used to resolve client metadata.
// It is internal and must not be exposed in model-list responses.
MetadataModelID string `json:"-"`
// ExplicitThinking indicates thinking/reasoning configuration was explicitly configured for this model.
ExplicitThinking bool `json:"-"`
// ExplicitInputModalities indicates input modalities were explicitly configured for this model.
ExplicitInputModalities bool `json:"-"`
// Object type for the model (typically "model")
Object string `json:"object"`
// Created timestamp when the model was created
Created int64 `json:"created"`
// OwnedBy indicates the organization that owns the model
OwnedBy string `json:"owned_by"`
// Type indicates the model type (e.g., "claude", "gemini", "openai")
Type string `json:"type"`
// DisplayName is the human-readable name for the model
DisplayName string `json:"display_name,omitempty"`
// Name is used for Gemini-style model names
Name string `json:"name,omitempty"`
// Version is the model version
Version string `json:"version,omitempty"`
// Description provides detailed information about the model
Description string `json:"description,omitempty"`
// InputTokenLimit is the maximum input token limit
InputTokenLimit int `json:"inputTokenLimit,omitempty"`
// OutputTokenLimit is the maximum output token limit
OutputTokenLimit int `json:"outputTokenLimit,omitempty"`
// SupportedGenerationMethods lists supported generation methods
SupportedGenerationMethods []string `json:"supportedGenerationMethods,omitempty"`
// ContextLength is the context window size
ContextLength int `json:"context_length,omitempty"`
// MaxContextLength is an explicit per-model context window override from configuration.
// It is carried internally for Codex client model catalog generation.
MaxContextLength int `json:"-"`
// MaxCompletionTokens is the maximum completion tokens
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
// SupportedParameters lists supported parameters
SupportedParameters []string `json:"supported_parameters,omitempty"`
// SupportedInputModalities lists supported input modalities (e.g., TEXT, IMAGE, VIDEO, AUDIO)
SupportedInputModalities []string `json:"supportedInputModalities,omitempty"`
// SupportedOutputModalities lists supported output modalities (e.g., TEXT, IMAGE)
SupportedOutputModalities []string `json:"supportedOutputModalities,omitempty"`
// SupportsWebSearch indicates this Antigravity model is listed by
// fetchAvailableModels.webSearchModelIds and can execute native googleSearch.
SupportsWebSearch bool `json:"supports_web_search,omitempty"`
// SupportConfigurationUpdate reports internal support for configuration_update.
SupportConfigurationUpdate bool `json:"-"`
// NativeCapabilities contains internal, static per-model capability metadata.
// It is intentionally separate from Antigravity's dynamically probed capability.
NativeCapabilities *NativeCapabilities `json:"-"`
// Thinking holds provider-specific reasoning/thinking budget capabilities.
// This is optional and currently used for Gemini thinking budget normalization.
Thinking *ThinkingSupport `json:"thinking,omitempty"`
// Config holds model-specific runtime overrides loaded from models.json.
Config *ModelConfig `json:"config,omitempty"`
// UserDefined indicates this model was defined through config file's models[]
// array (e.g., openai-compatibility.*.models[], *-api-key.models[]).
// UserDefined models have thinking configuration passed through without validation.
UserDefined bool `json:"-"`
// IsCompat enables compatibility handling for this configured API-key model.
// It is internal metadata and is not exposed in model listings.
IsCompat bool `json:"-"`
}
// ModelConfig holds optional runtime overrides for a model definition.
type ModelConfig struct {
// OverrideHeader forces upstream request headers when non-empty.
// Keys are header names (e.g. "user-agent"); values replace any existing header.
OverrideHeader map[string]string `json:"override_header,omitempty"`
}
// UnmarshalJSON loads internal capability metadata without exposing it
// through ModelInfo's normal JSON serialization.
func (m *ModelInfo) UnmarshalJSON(data []byte) error {
type modelInfoAlias ModelInfo
aux := struct {
*modelInfoAlias
NativeCapabilities *NativeCapabilities `json:"native_capabilities"`
SupportConfigurationUpdate bool `json:"support_configuration_update"`
}{modelInfoAlias: (*modelInfoAlias)(m)}
if errUnmarshal := json.Unmarshal(data, &aux); errUnmarshal != nil {
return errUnmarshal
}
m.NativeCapabilities = aux.NativeCapabilities
m.SupportConfigurationUpdate = aux.SupportConfigurationUpdate
return nil
}
type availableModelsCacheEntry struct {
models []map[string]any
expiresAt time.Time
}
// ThinkingSupport describes a model family's supported internal reasoning budget range.
// Values are interpreted in provider-native token units.
type ThinkingSupport struct {
// Min is the minimum allowed thinking budget (inclusive).
Min int `json:"min,omitempty" yaml:"min,omitempty"`
// Max is the maximum allowed thinking budget (inclusive).
Max int `json:"max,omitempty" yaml:"max,omitempty"`
// ZeroAllowed indicates whether 0 is a valid value (to disable thinking).
ZeroAllowed bool `json:"zero_allowed,omitempty" yaml:"zero-allowed,omitempty"`
// DynamicAllowed indicates whether -1 is a valid value (dynamic thinking budget).
DynamicAllowed bool `json:"dynamic_allowed,omitempty" yaml:"dynamic-allowed,omitempty"`
// DisableType is the Claude thinking.type sent when thinking is turned off.
// Empty means "disabled"; claude-sonnet-5-5 needs "between_tools".
DisableType string `json:"disable_type,omitempty" yaml:"disable-type,omitempty"`
// Levels defines discrete reasoning effort levels (e.g., "low", "medium", "high").
// When set, the model uses level-based reasoning instead of token budgets.
Levels []string `json:"levels,omitempty" yaml:"levels,omitempty"`
}
// ModelRegistration tracks a model's availability
type ModelRegistration struct {
// Info contains the model metadata
Info *ModelInfo
// InfoByProvider maps provider identifiers to specific ModelInfo to support differing capabilities.
InfoByProvider map[string]*ModelInfo
// Count is the number of active clients that can provide this model
Count int
// LastUpdated tracks when this registration was last modified
LastUpdated time.Time
// QuotaExceededClients tracks which clients have exceeded quota for this model
QuotaExceededClients map[string]*time.Time
// Providers tracks available clients grouped by provider identifier
Providers map[string]int
// SuspendedClients tracks temporarily disabled clients keyed by client ID
SuspendedClients map[string]string
// SuspendedClientDeadlines records optional expiry deadlines for suspended clients.
// Clients without an entry (or with a zero deadline) stay suspended until resumed.
SuspendedClientDeadlines map[string]time.Time
}
// ModelRegistryHook provides optional callbacks for external integrations to track model list changes.
// Hook implementations must be non-blocking and resilient; calls are executed asynchronously and panics are recovered.
type ModelRegistryHook interface {
OnModelsRegistered(ctx context.Context, provider, clientID string, models []*ModelInfo)
OnModelsUnregistered(ctx context.Context, provider, clientID string)
}
// ModelRegistry manages the global registry of available models
type ModelRegistry struct {
// models maps model ID to registration information
models map[string]*ModelRegistration
// clientModels maps client ID to the models it provides
clientModels map[string][]string
// clientModelInfos maps client ID to a map of model ID -> ModelInfo
// This preserves the original model info provided by each client
clientModelInfos map[string]map[string]*ModelInfo
// clientProviders maps client ID to its provider identifier
clientProviders map[string]string
// clientEpochs tracks monotonic registration epochs for each client ID
clientEpochs map[string]uint64
// clientGenerations tracks the latest generation applied for each client ID
clientGenerations map[string]uint64
// mutex ensures thread-safe access to the registry
mutex *sync.RWMutex
// availableModelsCache stores per-handler snapshots for GetAvailableModels.
availableModelsCache map[string]availableModelsCacheEntry
// webSearchClientCounts maps model ID to the number of clients advertising
// web search for it. It makes hasClientSupportingWebSearchLocked O(1)
// instead of scanning every client, which is quadratic when the whole
// registry is re-registered (thousands of clients x models).
webSearchClientCounts map[string]int
// webSearchProviderCounts maps model ID to provider -> client count for
// clients advertising web search, mirroring webSearchClientCounts.
webSearchProviderCounts map[string]map[string]int
// generation tracks changes to model registrations and availability.
generation uint64
// registrationEpoch tracks monotonic client registration and deregistration structural changes.
registrationEpoch atomic.Uint64
// hook is an optional callback sink for model registration changes
hook ModelRegistryHook
// nowFunc overrides time.Now when set (used by deterministic tests).
nowFunc func() time.Time
}
// Global model registry instance
var globalRegistry *ModelRegistry
var registryOnce sync.Once
// GetGlobalRegistry returns the global model registry instance
func GetGlobalRegistry() *ModelRegistry {
registryOnce.Do(func() {
globalRegistry = &ModelRegistry{
models: make(map[string]*ModelRegistration),
clientModels: make(map[string][]string),
clientModelInfos: make(map[string]map[string]*ModelInfo),
clientProviders: make(map[string]string),
clientEpochs: make(map[string]uint64),
clientGenerations: make(map[string]uint64),
availableModelsCache: make(map[string]availableModelsCacheEntry),
webSearchClientCounts: map[string]int{},
mutex: &sync.RWMutex{},
}
})
return globalRegistry
}
func (r *ModelRegistry) now() time.Time {
if r != nil && r.nowFunc != nil {
return r.nowFunc()
}
return time.Now()
}
func (r *ModelRegistry) ensureAvailableModelsCacheLocked() {
if r.availableModelsCache == nil {
r.availableModelsCache = make(map[string]availableModelsCacheEntry)
}
}
func (r *ModelRegistry) invalidateAvailableModelsCacheLocked() {
r.generation++
if len(r.availableModelsCache) == 0 {
return
}
clear(r.availableModelsCache)
}
// GetGeneration returns the current generation counter of model registrations.
func (r *ModelRegistry) GetGeneration() uint64 {
r.mutex.RLock()
defer r.mutex.RUnlock()
return r.generation
}
// RegistrationEpoch returns a monotonically increasing epoch that increments whenever
// client model registrations or deregistrations occur.
func (r *ModelRegistry) RegistrationEpoch() uint64 {
if r == nil {
return 0
}
return r.registrationEpoch.Load()
}
// LookupModelInfo searches dynamic registry (provider-specific > global) then static definitions.
func LookupModelInfo(modelID string, provider ...string) *ModelInfo {
modelID = strings.TrimSpace(modelID)
if modelID == "" {
return nil
}
p := ""
if len(provider) > 0 {
p = strings.ToLower(strings.TrimSpace(provider[0]))
}
if info := GetGlobalRegistry().GetModelInfo(modelID, p); info != nil {
return cloneModelInfo(info)
}
return cloneModelInfo(LookupStaticModelInfo(modelID))
}
// NativeCapabilityRoute describes one registered route to a public model.
type NativeCapabilityRoute struct {
Provider string
NativeCapabilities *NativeCapabilities
}
// ResolveResponsesWebSearchCapability conservatively resolves native web search
// across every route that can serve a public model. A known unsupported route or
// explicit model-level false wins; missing/unknown data produces unknown.
func ResolveResponsesWebSearchCapability(routes []NativeCapabilityRoute) *bool {
if len(routes) == 0 {
return nil
}
hasUnknown := false
for _, route := range routes {
if route.NativeCapabilities != nil && route.NativeCapabilities.WebSearch != nil && !*route.NativeCapabilities.WebSearch {
return boolPointer(false)
}
pathSupport := responsesWebSearchProviderPathSupport(route.Provider)
if pathSupport == nil {
hasUnknown = true
continue
}
if !*pathSupport {
return boolPointer(false)
}
if route.NativeCapabilities == nil || route.NativeCapabilities.WebSearch == nil {
hasUnknown = true
}
}
if hasUnknown {
return nil
}
return boolPointer(true)
}
func responsesWebSearchProviderPathSupport(provider string) *bool {
provider = strings.ToLower(strings.TrimSpace(provider))
switch provider {
case "codex", "xai", "claude", "antigravity":
return boolPointer(true)
case "openai", "openai-compatibility", "gemini", "aistudio", "vertex", "kimi", "kimi-ai", "kimi.ai", "kimi.com", "interactions", "gemini-interactions":
return boolPointer(false)
default:
if strings.HasPrefix(provider, "openai-compatible-") {
return boolPointer(false)
}
return nil
}
}
func boolPointer(value bool) *bool {
return &value
}
// GetResponsesWebSearchCapability resolves capability metadata across every
// registered client route for the exact public model ID.
func (r *ModelRegistry) GetResponsesWebSearchCapability(modelID string) *bool {
modelID = strings.TrimSpace(modelID)
if r == nil || modelID == "" {
return nil
}
r.mutex.RLock()
defer r.mutex.RUnlock()
routes := make([]NativeCapabilityRoute, 0)
for clientID, modelIDs := range r.clientModels {
for _, registeredID := range modelIDs {
if strings.TrimSpace(registeredID) != modelID {
continue
}
var nativeCapabilities *NativeCapabilities
if info := r.clientModelInfos[clientID][registeredID]; info != nil {
nativeCapabilities = info.NativeCapabilities
}
routes = append(routes, NativeCapabilityRoute{
Provider: r.clientProviders[clientID],
NativeCapabilities: nativeCapabilities,
})
}
}
return ResolveResponsesWebSearchCapability(routes)
}
// ModelOverrideHeaders returns models.json config.override_header for the model, if any.
// The returned map is a defensive copy and may be empty but never nil when overrides exist.
func ModelOverrideHeaders(modelID string, provider ...string) map[string]string {
info := LookupModelInfo(modelID, provider...)
if info == nil || info.Config == nil || len(info.Config.OverrideHeader) == 0 {
return nil
}
out := make(map[string]string, len(info.Config.OverrideHeader))
for key, value := range info.Config.OverrideHeader {
key = strings.TrimSpace(key)
if key == "" {
continue
}
out[key] = value
}
if len(out) == 0 {
return nil
}
return out
}
// SetHook sets an optional hook for observing model registration changes.
func (r *ModelRegistry) SetHook(hook ModelRegistryHook) {
if r == nil {
return
}
r.mutex.Lock()
defer r.mutex.Unlock()
r.hook = hook
}
const defaultModelRegistryHookTimeout = 5 * time.Second
const modelQuotaExceededWindow = 5 * time.Minute
func (r *ModelRegistry) triggerModelsRegistered(provider, clientID string, models []*ModelInfo) {
hook := r.hook
if hook == nil {
return
}
modelsCopy := cloneModelInfosUnique(models)
go func() {
defer func() {
if recovered := recover(); recovered != nil {
log.Errorf("model registry hook OnModelsRegistered panic: %v", recovered)
}
}()
ctx, cancel := context.WithTimeout(context.Background(), defaultModelRegistryHookTimeout)
defer cancel()
hook.OnModelsRegistered(ctx, provider, clientID, modelsCopy)
}()
}
func (r *ModelRegistry) triggerModelsUnregistered(provider, clientID string) {
hook := r.hook
if hook == nil {
return
}
go func() {
defer func() {
if recovered := recover(); recovered != nil {
log.Errorf("model registry hook OnModelsUnregistered panic: %v", recovered)
}
}()
ctx, cancel := context.WithTimeout(context.Background(), defaultModelRegistryHookTimeout)
defer cancel()
hook.OnModelsUnregistered(ctx, provider, clientID)
}()
}
// RegisterClient registers a client and its supported models
// Parameters:
// - clientID: Unique identifier for the client
// - clientProvider: Provider name (e.g., "gemini", "claude", "openai")
// - models: List of models that this client can provide
func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models []*ModelInfo) {
r.mutex.Lock()
defer r.mutex.Unlock()
r.registerClientLocked(clientID, clientProvider, models, false)
}
// RegisterClientIfEpochCurrent registers a client's models only if its
// registration epoch still equals expectedEpoch and publish is non-nil. The
// fence evaluation and the registration runs atomically under the registry
// mutex. publish must be fast and must not call back into the registry or
// block. The returned value reports whether the registration was applied.
func (r *ModelRegistry) RegisterClientIfEpochCurrent(clientID, clientProvider string, expectedEpoch uint64, models []*ModelInfo, publish func() bool) bool {
r.mutex.Lock()
defer r.mutex.Unlock()
if r.clientEpochs[clientID] != expectedEpoch || publish == nil || !publish() {
return false
}
r.registerClientLocked(clientID, clientProvider, models, false)
return true
}
// UnregisterClientIfEpochCurrent removes a client only if its registration
// epoch still equals expectedEpoch and publish is non-nil. The fence
// evaluation and the removal runs atomically under the registry mutex.
// publish must be fast and must not call back into the registry or block.
// The returned value reports whether the removal was applied.
func (r *ModelRegistry) UnregisterClientIfEpochCurrent(clientID string, expectedEpoch uint64, publish func() bool) bool {
r.mutex.Lock()
defer r.mutex.Unlock()
if r.clientEpochs[clientID] != expectedEpoch || publish == nil || !publish() {
return false
}
r.unregisterClientInternal(clientID)
r.invalidateAvailableModelsCacheLocked()
return true
}
// ReplaceClientModels replaces a client's models only if its registration epoch matches.
// Epoch zero and epochs for empty catalogs are valid. A successful replacement advances
// the epoch and preserves quota, suspension state and the projection generation
// watermark for an existing registration with the same provider.
// Empty models unregister the client; a mismatched epoch leaves the registry unchanged.
// The successful epoch is returned under the same lock as the replacement.
func (r *ModelRegistry) ReplaceClientModels(clientID, clientProvider string, expectedEpoch uint64, models []*ModelInfo) (uint64, bool) {
r.mutex.Lock()
defer r.mutex.Unlock()
if r.clientEpochs[clientID] != expectedEpoch {
return 0, false
}
r.registerClientLocked(clientID, clientProvider, models, true)
return r.clientEpochs[clientID], true
}
// registerClientLocked reconciles a client's models while the registry mutex is held.
func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, models []*ModelInfo, preserveState bool) {
r.ensureAvailableModelsCacheLocked()
if r.clientGenerations == nil {
r.clientGenerations = make(map[string]uint64)
}
if r.clientEpochs == nil {
r.clientEpochs = make(map[string]uint64)
}
provider := strings.ToLower(clientProvider)
uniqueModelIDs := make([]string, 0, len(models))
rawModelIDs := make([]string, 0, len(models))
newModels := make(map[string]*ModelInfo, len(models))
newCounts := make(map[string]int, len(models))
for _, model := range models {
if model == nil || model.ID == "" {
continue
}
rawModelIDs = append(rawModelIDs, model.ID)
newCounts[model.ID]++
if _, exists := newModels[model.ID]; exists {
continue
}
newModels[model.ID] = model
uniqueModelIDs = append(uniqueModelIDs, model.ID)
}
if len(uniqueModelIDs) == 0 {
// No models supplied; unregister existing client state if present.
// unregisterClientInternal retracts the client's web-search counts.
r.unregisterClientInternal(clientID)
delete(r.clientModels, clientID)
delete(r.clientModelInfos, clientID)
delete(r.clientProviders, clientID)
r.invalidateAvailableModelsCacheLocked()
misc.LogCredentialSeparator()
return
}
oldModels, hadExisting := r.clientModels[clientID]
oldProvider := r.clientProviders[clientID]
providerChanged := oldProvider != provider
// Retaining scheduling state also retains its projection watermark. An older
// request result may read the new epoch after capturing its auth snapshot.
r.clientEpochs[clientID]++
if !preserveState || !hadExisting || providerChanged {
r.clientGenerations[clientID] = uint64(0)
}
r.registrationEpoch.Add(1)
now := r.now()
if !hadExisting {
// Pure addition normally has no prior metadata. Clear orphaned metadata
// and refresh its models before publishing the new catalog.
r.clearOrphanedClientModelInfosLocked(clientID)
for _, modelID := range rawModelIDs {
model := newModels[modelID]
r.addModelRegistration(modelID, provider, model, now, clientID)
}
r.clientModels[clientID] = append([]string(nil), rawModelIDs...)
// Store client's own model infos
clientInfos := make(map[string]*ModelInfo, len(newModels))
for id, m := range newModels {
clientInfos[id] = cloneModelInfo(m)
}
r.clientModelInfos[clientID] = clientInfos
if provider != "" {
r.clientProviders[clientID] = provider
} else {
delete(r.clientProviders, clientID)
}
r.addClientWebSearchCountsLocked(clientID)
r.invalidateAvailableModelsCacheLocked()
r.triggerModelsRegistered(provider, clientID, models)
log.Debugf("Registered client %s from provider %s with %d models", clientID, clientProvider, len(rawModelIDs))
misc.LogCredentialSeparator()
return
}
oldCounts := make(map[string]int, len(oldModels))
for _, id := range oldModels {
oldCounts[id]++
}
added := make([]string, 0)
for _, id := range uniqueModelIDs {
if oldCounts[id] == 0 {
added = append(added, id)
}
}
removed := make([]string, 0)
for id := range oldCounts {
if newCounts[id] == 0 {
removed = append(removed, id)
}
}
// Handle provider change for overlapping models before modifications.
if providerChanged && oldProvider != "" {
for id, newCount := range newCounts {
if newCount == 0 {
continue
}
oldCount := oldCounts[id]
if oldCount == 0 {
continue
}
toRemove := newCount
if oldCount < toRemove {
toRemove = oldCount
}
if reg, ok := r.models[id]; ok && reg.Providers != nil {
if count, okProv := reg.Providers[oldProvider]; okProv {
if count <= toRemove {
delete(reg.Providers, oldProvider)
if reg.InfoByProvider != nil {
delete(reg.InfoByProvider, oldProvider)
}
} else {
reg.Providers[oldProvider] = count - toRemove
if reg.InfoByProvider != nil && reg.InfoByProvider[oldProvider] != nil {
reg.InfoByProvider[oldProvider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(id, oldProvider, clientID)
}
}
}
}
}
}
// Apply removals first to keep counters accurate.
for _, id := range removed {
oldCount := oldCounts[id]
for i := 0; i < oldCount; i++ {
r.removeModelRegistration(clientID, id, oldProvider, now, false)
}
}
for id, oldCount := range oldCounts {
newCount := newCounts[id]
if newCount == 0 || oldCount <= newCount {
continue
}
overage := oldCount - newCount
for i := 0; i < overage; i++ {
r.removeModelRegistration(clientID, id, oldProvider, now, preserveState && !providerChanged)
}
}
// Apply additions.
for id, newCount := range newCounts {
oldCount := oldCounts[id]
if newCount <= oldCount {
continue
}
model := newModels[id]
diff := newCount - oldCount
for i := 0; i < diff; i++ {
r.addModelRegistration(id, provider, model, now, clientID)
}
}
// Update metadata for models that remain associated with the client.
addedSet := make(map[string]struct{}, len(added))
for _, id := range added {
addedSet[id] = struct{}{}
}
for _, id := range uniqueModelIDs {
model := newModels[id]
if reg, ok := r.models[id]; ok {
hasWebSearch := model.SupportsWebSearch || r.hasClientSupportingWebSearchLocked(id, "", clientID)
reg.Info = cloneModelInfo(model)
reg.Info.SupportsWebSearch = hasWebSearch
if provider != "" {
if reg.InfoByProvider == nil {
reg.InfoByProvider = make(map[string]*ModelInfo)
}
hasProvWebSearch := model.SupportsWebSearch || r.hasClientSupportingWebSearchLocked(id, provider, clientID)
reg.InfoByProvider[provider] = cloneModelInfo(model)
reg.InfoByProvider[provider].SupportsWebSearch = hasProvWebSearch
}
if providerChanged && oldProvider != "" && reg.InfoByProvider != nil && reg.InfoByProvider[oldProvider] != nil {
reg.InfoByProvider[oldProvider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(id, oldProvider, clientID)
}
reg.LastUpdated = now
// Only conditional replacements preserve scheduling state for retained
// bindings with the same provider. Ordinary registration starts fresh.
if !preserveState || providerChanged || oldCounts[id] == 0 {
delete(reg.QuotaExceededClients, clientID)
delete(reg.SuspendedClients, clientID)
if reg.SuspendedClientDeadlines != nil {
delete(reg.SuspendedClientDeadlines, clientID)
}
}
if providerChanged && provider != "" {
if _, newlyAdded := addedSet[id]; newlyAdded {
continue
}
overlapCount := newCounts[id]
if oldCount := oldCounts[id]; oldCount < overlapCount {
overlapCount = oldCount
}
if overlapCount <= 0 {
continue
}
if reg.Providers == nil {
reg.Providers = make(map[string]int)
}
reg.Providers[provider] += overlapCount
}
}
}
// Update client bookkeeping.
// Retract this client's previous web-search contributions before rewriting
// clientModelInfos/clientProviders, then re-add them from the new state.
r.dropClientWebSearchCountsLocked(clientID)
if len(rawModelIDs) > 0 {
r.clientModels[clientID] = append([]string(nil), rawModelIDs...)
}
// Update client's own model infos
clientInfos := make(map[string]*ModelInfo, len(newModels))
for id, m := range newModels {
clientInfos[id] = cloneModelInfo(m)
}
r.clientModelInfos[clientID] = clientInfos
if provider != "" {
r.clientProviders[clientID] = provider
} else {
delete(r.clientProviders, clientID)
}
r.addClientWebSearchCountsLocked(clientID)
r.invalidateAvailableModelsCacheLocked()
r.triggerModelsRegistered(provider, clientID, models)
if len(added) == 0 && len(removed) == 0 && !providerChanged {
// Only metadata (e.g., display name) changed; keep no-op re-registration quiet.
return
}
log.Debugf("Reconciled client %s (provider %s) models: +%d, -%d", clientID, provider, len(added), len(removed))
misc.LogCredentialSeparator()
}
func (r *ModelRegistry) addModelRegistration(modelID, provider string, model *ModelInfo, now time.Time, excludeClientID string) {
if model == nil || modelID == "" {
return
}
if existing, exists := r.models[modelID]; exists {
existing.Count++
existing.LastUpdated = now
hasWebSearch := model.SupportsWebSearch || r.hasClientSupportingWebSearchLocked(modelID, "", excludeClientID)
existing.Info = cloneModelInfo(model)
existing.Info.SupportsWebSearch = hasWebSearch
if existing.SuspendedClients == nil {
existing.SuspendedClients = make(map[string]string)
}
if existing.SuspendedClientDeadlines == nil {
existing.SuspendedClientDeadlines = make(map[string]time.Time)
}
if existing.InfoByProvider == nil {
existing.InfoByProvider = make(map[string]*ModelInfo)
}
if provider != "" {
if existing.Providers == nil {
existing.Providers = make(map[string]int)
}
existing.Providers[provider]++
hasProvWebSearch := model.SupportsWebSearch || r.hasClientSupportingWebSearchLocked(modelID, provider, excludeClientID)
existing.InfoByProvider[provider] = cloneModelInfo(model)
existing.InfoByProvider[provider].SupportsWebSearch = hasProvWebSearch
}
log.Debugf("Incremented count for model %s, now %d clients", modelID, existing.Count)
return
}
registration := &ModelRegistration{
Info: cloneModelInfo(model),
InfoByProvider: make(map[string]*ModelInfo),
Count: 1,
LastUpdated: now,
QuotaExceededClients: make(map[string]*time.Time),
SuspendedClients: make(map[string]string),
SuspendedClientDeadlines: make(map[string]time.Time),
}
if provider != "" {
registration.Providers = map[string]int{provider: 1}
registration.InfoByProvider[provider] = cloneModelInfo(model)
}
r.models[modelID] = registration
log.Debugf("Registered new model %s from provider %s", modelID, provider)
}
// removeModelRegistration may preserve state when only a duplicate binding is removed.
func (r *ModelRegistry) removeModelRegistration(clientID, modelID, provider string, now time.Time, preserveState bool) {
registration, exists := r.models[modelID]
if !exists {
return
}
registration.Count--
registration.LastUpdated = now
if !preserveState {
delete(registration.QuotaExceededClients, clientID)
delete(registration.SuspendedClients, clientID)
if registration.SuspendedClientDeadlines != nil {
delete(registration.SuspendedClientDeadlines, clientID)
}
}
if registration.Count < 0 {
registration.Count = 0
}
if provider != "" && registration.Providers != nil {
if count, ok := registration.Providers[provider]; ok {
if count <= 1 {
delete(registration.Providers, provider)
if registration.InfoByProvider != nil {
delete(registration.InfoByProvider, provider)
}
} else {
registration.Providers[provider] = count - 1
}
}
}
log.Debugf("Decremented count for model %s, now %d clients", modelID, registration.Count)
if registration.Count <= 0 {
delete(r.models, modelID)
log.Debugf("Removed model %s as no clients remain", modelID)
} else {
if registration.Info != nil {
registration.Info.SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, "", clientID)
}
if provider != "" && registration.InfoByProvider != nil && registration.InfoByProvider[provider] != nil {
registration.InfoByProvider[provider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, provider, clientID)
}
}
}
func cloneModelInfo(model *ModelInfo) *ModelInfo {
if model == nil {
return nil
}
copyModel := *model
if model.NativeCapabilities != nil {
copyCapabilities := *model.NativeCapabilities
if model.NativeCapabilities.WebSearch != nil {
webSearch := *model.NativeCapabilities.WebSearch
copyCapabilities.WebSearch = &webSearch
}
copyModel.NativeCapabilities = ©Capabilities
}
if len(model.SupportedGenerationMethods) > 0 {
copyModel.SupportedGenerationMethods = append([]string(nil), model.SupportedGenerationMethods...)
}
if len(model.SupportedParameters) > 0 {
copyModel.SupportedParameters = append([]string(nil), model.SupportedParameters...)
}
if len(model.SupportedInputModalities) > 0 {
copyModel.SupportedInputModalities = append([]string(nil), model.SupportedInputModalities...)
}
if len(model.SupportedOutputModalities) > 0 {
copyModel.SupportedOutputModalities = append([]string(nil), model.SupportedOutputModalities...)
}
if model.Thinking != nil {
copyThinking := *model.Thinking
if len(model.Thinking.Levels) > 0 {
copyThinking.Levels = append([]string(nil), model.Thinking.Levels...)
}
copyModel.Thinking = ©Thinking
}
if model.Config != nil {
copyConfig := *model.Config
if len(model.Config.OverrideHeader) > 0 {
copyConfig.OverrideHeader = make(map[string]string, len(model.Config.OverrideHeader))
for key, value := range model.Config.OverrideHeader {
copyConfig.OverrideHeader[key] = value
}
}
copyModel.Config = ©Config
}
return ©Model
}
func cloneModelInfosUnique(models []*ModelInfo) []*ModelInfo {
if len(models) == 0 {
return nil
}
cloned := make([]*ModelInfo, 0, len(models))
seen := make(map[string]struct{}, len(models))
for _, model := range models {
if model == nil || model.ID == "" {
continue
}
if _, exists := seen[model.ID]; exists {
continue
}
seen[model.ID] = struct{}{}
cloned = append(cloned, cloneModelInfo(model))
}
return cloned
}
// UnregisterClient removes a client and decrements counts for its models
// Parameters:
// - clientID: Unique identifier for the client to remove
func (r *ModelRegistry) UnregisterClient(clientID string) {
r.mutex.Lock()
defer r.mutex.Unlock()
r.unregisterClientInternal(clientID)
r.invalidateAvailableModelsCacheLocked()
}
// unregisterClientInternal performs the actual client unregistration (internal, no locking)
func (r *ModelRegistry) unregisterClientInternal(clientID string) {
if r.clientGenerations == nil {
r.clientGenerations = make(map[string]uint64)
}
if r.clientEpochs == nil {
r.clientEpochs = make(map[string]uint64)
}
r.clientEpochs[clientID]++
r.clientGenerations[clientID]++
r.registrationEpoch.Add(1)
models, exists := r.clientModels[clientID]
provider, hasProvider := r.clientProviders[clientID]
if !exists {
r.clearOrphanedClientModelInfosLocked(clientID)
if hasProvider {
delete(r.clientProviders, clientID)
}
return
}
now := r.now()
for _, modelID := range models {
if registration, isExists := r.models[modelID]; isExists {
registration.Count--
registration.LastUpdated = now
// Remove quota tracking for this client
delete(registration.QuotaExceededClients, clientID)
if registration.SuspendedClients != nil {
delete(registration.SuspendedClients, clientID)
}
if registration.SuspendedClientDeadlines != nil {
delete(registration.SuspendedClientDeadlines, clientID)
}
if hasProvider && registration.Providers != nil {
if count, ok := registration.Providers[provider]; ok {
if count <= 1 {
delete(registration.Providers, provider)
if registration.InfoByProvider != nil {
delete(registration.InfoByProvider, provider)
}
} else {
registration.Providers[provider] = count - 1
}
}
}
log.Debugf("Decremented count for model %s, now %d clients", modelID, registration.Count)
// Remove model if no clients remain
if registration.Count <= 0 {
delete(r.models, modelID)
log.Debugf("Removed model %s as no clients remain", modelID)
} else {
if registration.Info != nil {
registration.Info.SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, "", clientID)
}
if hasProvider && registration.InfoByProvider != nil && registration.InfoByProvider[provider] != nil {
registration.InfoByProvider[provider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, provider, clientID)
}
}
}
}
// Retract this client's web-search contributions while clientModelInfos and
// clientProviders still describe it. This must run after the per-model loop
// above, which relies on excludeClientID to skip this client.
r.dropClientWebSearchCountsLocked(clientID)
delete(r.clientModels, clientID)
delete(r.clientModelInfos, clientID)
if hasProvider {
delete(r.clientProviders, clientID)
}
log.Debugf("Unregistered client %s", clientID)
// Separator line after completing client unregistration (after the summary line)
misc.LogCredentialSeparator()
r.triggerModelsUnregistered(provider, clientID)
}
// SetModelQuotaExceeded marks a model as quota exceeded for a specific client
// Parameters:
// - clientID: The client that exceeded quota
// - modelID: The model that exceeded quota
func (r *ModelRegistry) SetModelQuotaExceeded(clientID, modelID string) {
r.mutex.Lock()
defer r.mutex.Unlock()
r.ensureAvailableModelsCacheLocked()
if registration, exists := r.models[modelID]; exists {
now := r.now()
registration.QuotaExceededClients[clientID] = &now
r.invalidateAvailableModelsCacheLocked()
log.Debugf("Marked model %s as quota exceeded for client %s", modelID, clientID)
}
}
// ClearModelQuotaExceeded removes quota exceeded status for a model and client
// Parameters:
// - clientID: The client to clear quota status for
// - modelID: The model to clear quota status for
func (r *ModelRegistry) ClearModelQuotaExceeded(clientID, modelID string) {
r.mutex.Lock()
defer r.mutex.Unlock()
r.ensureAvailableModelsCacheLocked()
if registration, exists := r.models[modelID]; exists {
delete(registration.QuotaExceededClients, clientID)
r.invalidateAvailableModelsCacheLocked()
// log.Debugf("Cleared quota exceeded status for model %s and client %s", modelID, clientID)
}
}
// ClientModelProjection describes the desired availability and quota state for a client's model.
type ClientModelProjection struct {
ModelID string
Suspended bool
SuspendReason string
// SuspendUntil optionally bounds the suspension in time. Once it passes,
// the client counts as available again without requiring re-projection.
SuspendUntil time.Time
QuotaExceeded bool
}
// ApplyClientModelProjections atomically applies model suspension and quota exceeded state for clientID
// if the supplied epoch and generation are current. Stale or mismatched epoch (!= current client epoch) or stale generation
// (< current client generation), or projections for unregistered clients/models are rejected.
func (r *ModelRegistry) ApplyClientModelProjections(clientID string, epoch uint64, generation uint64, projections []ClientModelProjection) bool {
if r == nil {
return false
}
clientID = strings.TrimSpace(clientID)
if clientID == "" {
return false
}
r.mutex.Lock()
defer r.mutex.Unlock()
if r.clientEpochs == nil {
r.clientEpochs = make(map[string]uint64)
}
if r.clientGenerations == nil {
r.clientGenerations = make(map[string]uint64)
}
clientRegisteredModels, exists := r.clientModels[clientID]
if !exists || len(clientRegisteredModels) == 0 {
return false
}
currentEpoch := r.clientEpochs[clientID]
if epoch != currentEpoch {
return false
}
currentGen := r.clientGenerations[clientID]
if generation < currentGen {
return false
}
registeredSet := make(map[string]struct{}, len(clientRegisteredModels))
for _, m := range clientRegisteredModels {
trimmedModel := strings.TrimSpace(m)
if trimmedModel != "" {
registeredSet[trimmedModel] = struct{}{}
}
}
hasValidModel := false
for _, proj := range projections {
modelID := strings.TrimSpace(proj.ModelID)
if modelID == "" {
continue
}
if _, isOwned := registeredSet[modelID]; isOwned {
if registration, existsModel := r.models[modelID]; existsModel && registration != nil {
hasValidModel = true
break
}
}
}
if !hasValidModel {
return false
}
r.clientGenerations[clientID] = generation
r.ensureAvailableModelsCacheLocked()
now := r.now()
changed := false
for _, proj := range projections {
modelID := strings.TrimSpace(proj.ModelID)
if modelID == "" {
continue
}
if _, isOwned := registeredSet[modelID]; !isOwned {
continue
}
registration, existsModel := r.models[modelID]
if !existsModel || registration == nil {
continue
}
// Update suspended state
if proj.Suspended {
if registration.SuspendedClients == nil {
registration.SuspendedClients = make(map[string]string)
}
if registration.SuspendedClientDeadlines == nil {
registration.SuspendedClientDeadlines = make(map[string]time.Time)
}
deadlineChanged := false
if currentDeadline, ok := registration.SuspendedClientDeadlines[clientID]; ok {
deadlineChanged = !proj.SuspendUntil.Equal(currentDeadline)
} else {
deadlineChanged = !proj.SuspendUntil.IsZero()
}
if currentReason, ok := registration.SuspendedClients[clientID]; !ok || currentReason != proj.SuspendReason || deadlineChanged {
registration.SuspendedClients[clientID] = proj.SuspendReason
if proj.SuspendUntil.IsZero() {
delete(registration.SuspendedClientDeadlines, clientID)
} else {
registration.SuspendedClientDeadlines[clientID] = proj.SuspendUntil
}
registration.LastUpdated = now
changed = true
}
} else {
if registration.SuspendedClients != nil {
if _, ok := registration.SuspendedClients[clientID]; ok {
delete(registration.SuspendedClients, clientID)
if registration.SuspendedClientDeadlines != nil {
delete(registration.SuspendedClientDeadlines, clientID)
}
registration.LastUpdated = now
changed = true
}
}
}
// Update quota exceeded state
if proj.QuotaExceeded {
if registration.QuotaExceededClients == nil {
registration.QuotaExceededClients = make(map[string]*time.Time)
}
if _, ok := registration.QuotaExceededClients[clientID]; !ok {
registration.QuotaExceededClients[clientID] = &now
registration.LastUpdated = now
changed = true
}
} else {
if registration.QuotaExceededClients != nil {
if _, ok := registration.QuotaExceededClients[clientID]; ok {
delete(registration.QuotaExceededClients, clientID)
registration.LastUpdated = now
changed = true
}
}
}
}
if changed {
r.invalidateAvailableModelsCacheLocked()
}
return true
}
// ApplyClientModelCapabilities applies capability mutations to matching models of clientID
// if the client is currently registered and its registration epoch matches expectedEpoch.
// Returns true if applied, false if client is unregistered or epoch changed.
func (r *ModelRegistry) ApplyClientModelCapabilities(clientID string, expectedEpoch uint64, mutate func(modelID string, info *ModelInfo)) bool {
if r == nil || mutate == nil {
return false
}
clientID = strings.TrimSpace(clientID)
if clientID == "" {
return false
}
r.mutex.Lock()
defer r.mutex.Unlock()
if r.clientEpochs == nil || r.clientEpochs[clientID] != expectedEpoch {
return false
}
clientInfos, exists := r.clientModelInfos[clientID]
if !exists || len(clientInfos) == 0 {
return false
}
provider := r.clientProviders[clientID]
// The callback mutates capabilities in place, which can add or clear this
// client's web-search contributions. Retract the current ones first so the
// index stays in sync with clientInfos.
r.dropClientWebSearchCountsLocked(clientID)
for id, info := range clientInfos {
if info != nil {
mutate(id, info)
// Reinstate this client's contribution before querying the index so
// the aggregate reflects the post-mutation state, including itself.
if info.SupportsWebSearch {
r.incrementWebSearchCountLocked(id, provider)
}
if reg, okReg := r.models[id]; okReg && reg != nil {
hasWebSearch := r.hasClientSupportingWebSearchLocked(id, "", "")
if reg.Info != nil {
reg.Info.SupportsWebSearch = hasWebSearch
}
if provider != "" && reg.InfoByProvider != nil && reg.InfoByProvider[provider] != nil {
reg.InfoByProvider[provider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(id, provider, "")
}
}
}
}
r.invalidateAvailableModelsCacheLocked()
return true
}
// hasClientSupportingWebSearchLocked excludes the given ID even when it is empty,
// matching the original scan without walking the client set.
func (r *ModelRegistry) hasClientSupportingWebSearchLocked(modelID, provider, excludeClientID string) bool {
if provider != "" {
count := r.webSearchProviderCounts[modelID][provider]
if count <= 0 {
return false
}
if r.clientProviders[excludeClientID] == provider {
if info := r.clientModelInfos[excludeClientID][modelID]; info != nil && info.SupportsWebSearch {
count--
}
}
return count > 0
}
count := r.webSearchClientCounts[modelID]
if count <= 0 {
return false
}
if info := r.clientModelInfos[excludeClientID][modelID]; info != nil && info.SupportsWebSearch {
count--
}
return count > 0
}
// incrementWebSearchCountLocked records one client's web-search contribution for modelID.
func (r *ModelRegistry) incrementWebSearchCountLocked(modelID, provider string) {
if modelID == "" {
return
}
if r.webSearchClientCounts == nil {
r.webSearchClientCounts = make(map[string]int)
}
r.webSearchClientCounts[modelID]++
if provider == "" {
return
}
if r.webSearchProviderCounts == nil {
r.webSearchProviderCounts = make(map[string]map[string]int)
}
counts := r.webSearchProviderCounts[modelID]
if counts == nil {
counts = make(map[string]int)
r.webSearchProviderCounts[modelID] = counts
}
counts[provider]++
}
// decrementWebSearchCountLocked removes one client's web-search contribution for modelID.
func (r *ModelRegistry) decrementWebSearchCountLocked(modelID, provider string) {
if modelID == "" {
return
}
if count := r.webSearchClientCounts[modelID]; count > 1 {
r.webSearchClientCounts[modelID] = count - 1
} else {
delete(r.webSearchClientCounts, modelID)
}
if provider == "" {
return
}
counts := r.webSearchProviderCounts[modelID]
if counts == nil {
return
}
if count := counts[provider]; count > 1 {
counts[provider] = count - 1
} else {
delete(counts, provider)
}
if len(counts) == 0 {
delete(r.webSearchProviderCounts, modelID)
}
}
// clearOrphanedClientModelInfosLocked refreshes capabilities before retracting
// and discarding metadata for a client without clientModels. The caller must
// retain clientProviders until this returns; model multiplicities are unknown.
func (r *ModelRegistry) clearOrphanedClientModelInfosLocked(clientID string) {
provider, hasProvider := r.clientProviders[clientID]
for modelID := range r.clientModelInfos[clientID] {
if registration := r.models[modelID]; registration != nil {
if registration.Info != nil {
registration.Info.SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, "", clientID)
}
if hasProvider && registration.InfoByProvider[provider] != nil {
registration.InfoByProvider[provider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(modelID, provider, clientID)
}
}
}
r.dropClientWebSearchCountsLocked(clientID)
delete(r.clientModelInfos, clientID)
}
// dropClientWebSearchCountsLocked removes every web-search contribution currently
// recorded for clientID. It must run before clientModelInfos/clientProviders are
// rewritten for that client.
func (r *ModelRegistry) dropClientWebSearchCountsLocked(clientID string) {
infos := r.clientModelInfos[clientID]
if len(infos) == 0 {
return
}
provider := r.clientProviders[clientID]
for modelID, info := range infos {
if info != nil && info.SupportsWebSearch {
r.decrementWebSearchCountLocked(modelID, provider)
}
}
}
// addClientWebSearchCountsLocked records every web-search contribution currently
// present for clientID. It must run after clientModelInfos/clientProviders are
// rewritten for that client.
func (r *ModelRegistry) addClientWebSearchCountsLocked(clientID string) {
infos := r.clientModelInfos[clientID]
if len(infos) == 0 {
return
}
provider := r.clientProviders[clientID]
for modelID, info := range infos {
if info != nil && info.SupportsWebSearch {
r.incrementWebSearchCountLocked(modelID, provider)
}
}
}
// SuspendClientModel marks a client's model as temporarily unavailable until explicitly resumed.
// Parameters:
// - clientID: The client to suspend
// - modelID: The model affected by the suspension
// - reason: Optional description for observability
func (r *ModelRegistry) SuspendClientModel(clientID, modelID, reason string) {
if clientID == "" || modelID == "" {
return
}
r.mutex.Lock()
defer r.mutex.Unlock()
r.ensureAvailableModelsCacheLocked()
registration, exists := r.models[modelID]
if !exists || registration == nil {
return
}
if registration.SuspendedClients == nil {
registration.SuspendedClients = make(map[string]string)
}
now := r.now()
_, _, active := clientSuspensionActive(registration, clientID, now)
hadDeadline := false
if registration.SuspendedClientDeadlines != nil {
_, hadDeadline = registration.SuspendedClientDeadlines[clientID]
}
if active && !hadDeadline {
return
}
registration.SuspendedClients[clientID] = reason
if registration.SuspendedClientDeadlines != nil {
delete(registration.SuspendedClientDeadlines, clientID)
}
registration.LastUpdated = now
r.invalidateAvailableModelsCacheLocked()
if reason != "" {
log.Debugf("Suspended client %s for model %s: %s", clientID, modelID, reason)
} else {
log.Debugf("Suspended client %s for model %s", clientID, modelID)
}
}
// ResumeClientModel clears a previous suspension so the client counts toward availability again.
// Parameters:
// - clientID: The client to resume
// - modelID: The model being resumed
func (r *ModelRegistry) ResumeClientModel(clientID, modelID string) {
if clientID == "" || modelID == "" {
return
}
r.mutex.Lock()
defer r.mutex.Unlock()
r.ensureAvailableModelsCacheLocked()
registration, exists := r.models[modelID]
if !exists || registration == nil || registration.SuspendedClients == nil {
return
}
if _, ok := registration.SuspendedClients[clientID]; !ok {
return
}
delete(registration.SuspendedClients, clientID)
if registration.SuspendedClientDeadlines != nil {
delete(registration.SuspendedClientDeadlines, clientID)
}
registration.LastUpdated = r.now()
r.invalidateAvailableModelsCacheLocked()
log.Debugf("Resumed client %s for model %s", clientID, modelID)
}
// ClientSupportsModel reports whether the client registered support for modelID.
func (r *ModelRegistry) ClientSupportsModel(clientID, modelID string) bool {
clientID = strings.TrimSpace(clientID)
modelID = strings.TrimSpace(modelID)
if clientID == "" || modelID == "" {
return false
}
r.mutex.RLock()
defer r.mutex.RUnlock()
models, exists := r.clientModels[clientID]
if !exists || len(models) == 0 {
return false
}
for _, id := range models {
if strings.EqualFold(strings.TrimSpace(id), modelID) {
return true
}
}
return false
}
// IsModelSuspendedForClient reports whether a model is currently suspended for a specific client.
func (r *ModelRegistry) IsModelSuspendedForClient(clientID, modelID string) bool {
clientID = strings.TrimSpace(clientID)
modelID = strings.TrimSpace(modelID)
if clientID == "" || modelID == "" {
return false
}
now := r.now()
r.mutex.RLock()
defer r.mutex.RUnlock()
registration, exists := r.models[modelID]
if !exists || registration == nil {
return false
}
_, _, active := clientSuspensionActive(registration, clientID, now)
return active
}
// IsModelQuotaExceededForClient reports whether a model is currently marked quota exceeded for a specific client.
func (r *ModelRegistry) IsModelQuotaExceededForClient(clientID, modelID string) bool {
clientID = strings.TrimSpace(clientID)
modelID = strings.TrimSpace(modelID)
if clientID == "" || modelID == "" {
return false
}
r.mutex.RLock()
defer r.mutex.RUnlock()
registration, exists := r.models[modelID]
if !exists || registration == nil || registration.QuotaExceededClients == nil {
return false
}
_, exceeded := registration.QuotaExceededClients[clientID]
return exceeded
}
// GetAvailableModels returns all models that have at least one available client
// Parameters:
// - handlerType: The handler type to filter models for (e.g., "openai", "claude", "gemini")
//
// Returns:
// - []map[string]any: List of available models in the requested format
func (r *ModelRegistry) GetAvailableModels(handlerType string) []map[string]any {
now := r.now()
r.mutex.RLock()
if cache, ok := r.availableModelsCache[handlerType]; ok && (cache.expiresAt.IsZero() || now.Before(cache.expiresAt)) {
models := cloneModelMaps(cache.models)
r.mutex.RUnlock()
return models
}
r.mutex.RUnlock()
r.mutex.Lock()
defer r.mutex.Unlock()
r.ensureAvailableModelsCacheLocked()
if cache, ok := r.availableModelsCache[handlerType]; ok && (cache.expiresAt.IsZero() || now.Before(cache.expiresAt)) {
return cloneModelMaps(cache.models)
}
models, expiresAt := r.buildAvailableModelsLocked(handlerType, now)
r.availableModelsCache[handlerType] = availableModelsCacheEntry{
models: cloneModelMaps(models),
expiresAt: expiresAt,
}
return models
}
// isQuotaSuspension reports whether a suspension is a quota cooldown. A model in a
// quota cooldown stays listed: the limit is temporary, and a request meanwhile gets a
// 429 with its retry time instead of an unknown model. A credential-wide quota
// ("credential_quota") suspends every model of that credential, so it is a quota
// cooldown too; counting it as another failure hides all of them when no other
// credential serves them.
func isQuotaSuspension(reason string) bool {
return strings.EqualFold(reason, "quota") || strings.EqualFold(reason, "credential_quota")
}
func clientSuspensionActive(registration *ModelRegistration, clientID string, now time.Time) (string, time.Time, bool) {
if registration == nil || len(registration.SuspendedClients) == 0 {
return "", time.Time{}, false
}
reason, ok := registration.SuspendedClients[clientID]
if !ok {
return "", time.Time{}, false
}
if deadline, hasDeadline := registration.SuspendedClientDeadlines[clientID]; hasDeadline && !deadline.IsZero() {
if !now.Before(deadline) {
return "", time.Time{}, false
}
return reason, deadline, true
}
return reason, time.Time{}, true
}
func clientQuotaRecoveryAt(registration *ModelRegistration, clientID string, quotaTime *time.Time) (time.Time, bool) {
if registration == nil || quotaTime == nil {
return time.Time{}, false
}
recoveryAt := quotaTime.Add(modelQuotaExceededWindow)
if deadline, hasDeadline := registration.SuspendedClientDeadlines[clientID]; hasDeadline && !deadline.IsZero() && deadline.Before(recoveryAt) {
recoveryAt = deadline
}
return recoveryAt, true
}
func modelRegistrationAvailability(registration *ModelRegistration, now time.Time) (bool, time.Time) {
if registration == nil {
return false, time.Time{}
}
availableClients := registration.Count
expiredClients := 0
var expiresAt time.Time
for clientID, quotaTime := range registration.QuotaExceededClients {
recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, quotaTime)
if !ok {
continue
}
if now.Before(recoveryAt) {
expiredClients++
if expiresAt.IsZero() || recoveryAt.Before(expiresAt) {
expiresAt = recoveryAt
}
}
}
cooldownSuspended := 0
otherSuspended := 0
quotaAndSuspended := 0
if registration.SuspendedClients != nil {
for clientID := range registration.SuspendedClients {
reason, deadline, active := clientSuspensionActive(registration, clientID, now)
if !active {
continue
}
if !deadline.IsZero() && (expiresAt.IsZero() || deadline.Before(expiresAt)) {
expiresAt = deadline
}
if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, registration.QuotaExceededClients[clientID]); ok && now.Before(recoveryAt) {
quotaAndSuspended++
}
if isQuotaSuspension(reason) {
cooldownSuspended++
continue
}
otherSuspended++
}
}
// A quota cooldown can mark the same client both quota-exceeded and
// suspended. Count that unavailable client only once.
effectiveClients := availableClients - expiredClients - cooldownSuspended - otherSuspended + quotaAndSuspended
if effectiveClients < 0 {
effectiveClients = 0
}
available := effectiveClients > 0 || (availableClients > 0 && (expiredClients > 0 || cooldownSuspended > 0) && otherSuspended == 0)
return available, expiresAt
}
// GetAvailableModelInfos returns cloned metadata for all currently available models.
func (r *ModelRegistry) GetAvailableModelInfos() []*ModelInfo {
now := r.now()
r.mutex.RLock()
defer r.mutex.RUnlock()
result := make([]*ModelInfo, 0, len(r.models))
for _, registration := range r.models {
available, _ := modelRegistrationAvailability(registration, now)
if !available || registration == nil || registration.Info == nil {
continue
}
result = append(result, cloneModelInfo(registration.Info))
}
sort.Slice(result, func(i, j int) bool {
return strings.TrimSpace(result[i].ID) < strings.TrimSpace(result[j].ID)
})
return result
}
func (r *ModelRegistry) buildAvailableModelsLocked(handlerType string, now time.Time) ([]map[string]any, time.Time) {
models := make([]map[string]any, 0, len(r.models))
var expiresAt time.Time
for _, registration := range r.models {
available, registrationExpiresAt := modelRegistrationAvailability(registration, now)
if !registrationExpiresAt.IsZero() && (expiresAt.IsZero() || registrationExpiresAt.Before(expiresAt)) {
expiresAt = registrationExpiresAt
}
if !available || registration == nil {
continue
}
model := r.convertModelToMap(registration.Info, handlerType)
if model != nil {
models = append(models, model)
}
}
return models, expiresAt
}
func cloneModelMaps(models []map[string]any) []map[string]any {
cloned := make([]map[string]any, 0, len(models))
for _, model := range models {
if model == nil {
cloned = append(cloned, nil)
continue
}
copyModel := make(map[string]any, len(model))
for key, value := range model {
copyModel[key] = cloneModelMapValue(value)
}
cloned = append(cloned, copyModel)
}
return cloned
}
func cloneModelMapValue(value any) any {
switch typed := value.(type) {
case map[string]any:
copyMap := make(map[string]any, len(typed))
for key, entry := range typed {
copyMap[key] = cloneModelMapValue(entry)
}
return copyMap
case []any:
copySlice := make([]any, len(typed))
for i, entry := range typed {
copySlice[i] = cloneModelMapValue(entry)
}
return copySlice
case []string:
return append([]string(nil), typed...)
default:
return value
}
}
// GetAvailableModelsByProvider returns models available for the given provider identifier.
// Parameters:
// - provider: Provider identifier (e.g., "codex", "gemini", "antigravity")
//
// Returns:
// - []*ModelInfo: List of available models for the provider
func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelInfo {
provider = strings.ToLower(strings.TrimSpace(provider))
if provider == "" {
return nil
}
r.mutex.RLock()
defer r.mutex.RUnlock()
type providerModel struct {
count int
info *ModelInfo
}
providerModels := make(map[string]*providerModel)
for clientID, clientProvider := range r.clientProviders {
if clientProvider != provider {
continue
}
modelIDs := r.clientModels[clientID]
if len(modelIDs) == 0 {
continue
}
clientInfos := r.clientModelInfos[clientID]
for _, modelID := range modelIDs {
modelID = strings.TrimSpace(modelID)
if modelID == "" {
continue
}
entry := providerModels[modelID]
if entry == nil {
entry = &providerModel{}
providerModels[modelID] = entry
}
entry.count++
if entry.info == nil {
if clientInfos != nil {
if info := clientInfos[modelID]; info != nil {
entry.info = info
}
}
if entry.info == nil {
if reg, ok := r.models[modelID]; ok && reg != nil && reg.Info != nil {
entry.info = reg.Info
}
}
}
}
}
if len(providerModels) == 0 {
return nil
}
now := r.now()
result := make([]*ModelInfo, 0, len(providerModels))
for modelID, entry := range providerModels {
if entry == nil || entry.count <= 0 {
continue
}
registration, ok := r.models[modelID]
expiredClients := 0
cooldownSuspended := 0
otherSuspended := 0
quotaAndSuspended := 0
if ok && registration != nil {
if registration.QuotaExceededClients != nil {
for clientID, quotaTime := range registration.QuotaExceededClients {
if clientID == "" {
continue
}
if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider {
continue
}
if recoveryAt, okRecovery := clientQuotaRecoveryAt(registration, clientID, quotaTime); okRecovery && now.Before(recoveryAt) {
expiredClients++
}
}
}
if registration.SuspendedClients != nil {
for clientID := range registration.SuspendedClients {
if clientID == "" {
continue
}
if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider {
continue
}
reason, _, active := clientSuspensionActive(registration, clientID, now)
if !active {
continue
}
if recoveryAt, okRecovery := clientQuotaRecoveryAt(registration, clientID, registration.QuotaExceededClients[clientID]); okRecovery && now.Before(recoveryAt) {
quotaAndSuspended++
}
if isQuotaSuspension(reason) {
cooldownSuspended++
continue
}
otherSuspended++
}
}
}
availableClients := entry.count
effectiveClients := availableClients - expiredClients - cooldownSuspended - otherSuspended + quotaAndSuspended
if effectiveClients < 0 {
effectiveClients = 0
}
if effectiveClients > 0 || (availableClients > 0 && (expiredClients > 0 || cooldownSuspended > 0) && otherSuspended == 0) {
if entry.info != nil {
result = append(result, cloneModelInfo(entry.info))
continue
}
if ok && registration != nil && registration.Info != nil {
result = append(result, cloneModelInfo(registration.Info))
}
}
}
return result
}
// GetModelCount returns the number of available clients for a specific model
// Parameters:
// - modelID: The model ID to check
//
// Returns:
// - int: Number of available clients for the model
func (r *ModelRegistry) GetModelCount(modelID string) int {
r.mutex.RLock()
defer r.mutex.RUnlock()
if registration, exists := r.models[modelID]; exists {
now := r.now()
// Count clients that have exceeded quota but haven't recovered yet
expiredClients := 0
for clientID, quotaTime := range registration.QuotaExceededClients {
if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, quotaTime); ok && now.Before(recoveryAt) {
expiredClients++
}
}
suspendedClients := 0
for clientID := range registration.SuspendedClients {
if _, _, active := clientSuspensionActive(registration, clientID, now); !active {
continue
}
if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, registration.QuotaExceededClients[clientID]); ok && now.Before(recoveryAt) {
continue
}
suspendedClients++
}
result := registration.Count - expiredClients - suspendedClients
if result < 0 {
return 0
}
return result
}
return 0
}
// GetModelProviders returns provider identifiers that currently supply the given model
// Parameters:
// - modelID: The model ID to check
//
// Returns:
// - []string: Provider identifiers ordered by availability count (descending)
func (r *ModelRegistry) GetModelProviders(modelID string) []string {
r.mutex.RLock()
defer r.mutex.RUnlock()
registration, exists := r.models[modelID]
if !exists || registration == nil || len(registration.Providers) == 0 {
return nil
}
type providerCount struct {
name string
count int
}
providers := make([]providerCount, 0, len(registration.Providers))
// suspendedByProvider := make(map[string]int)
// if registration.SuspendedClients != nil {
// for clientID := range registration.SuspendedClients {
// if provider, ok := r.clientProviders[clientID]; ok && provider != "" {
// suspendedByProvider[provider]++
// }
// }
// }
for name, count := range registration.Providers {
if count <= 0 {
continue
}
// adjusted := count - suspendedByProvider[name]
// if adjusted <= 0 {
// continue
// }
// providers = append(providers, providerCount{name: name, count: adjusted})
providers = append(providers, providerCount{name: name, count: count})
}
if len(providers) == 0 {
return nil
}
sort.Slice(providers, func(i, j int) bool {
if providers[i].count == providers[j].count {
return providers[i].name < providers[j].name
}
return providers[i].count > providers[j].count
})
result := make([]string, 0, len(providers))
for _, item := range providers {
result = append(result, item.name)
}
return result
}
// GetModelInfo returns ModelInfo, prioritizing provider-specific definition if available.
func (r *ModelRegistry) GetModelInfo(modelID, provider string) *ModelInfo {
r.mutex.RLock()
defer r.mutex.RUnlock()
if reg, ok := r.models[modelID]; ok && reg != nil {
// Try provider specific definition first
if provider != "" && reg.InfoByProvider != nil {
if reg.Providers != nil {
if count, ok := reg.Providers[provider]; ok && count > 0 {
if info, ok := reg.InfoByProvider[provider]; ok && info != nil {
return cloneModelInfo(info)
}
}
}
}
// Fallback to global info (last registered)
return cloneModelInfo(reg.Info)
}
return nil
}
// convertModelToMap converts ModelInfo to the appropriate format for different handler types
func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) map[string]any {
if model == nil {
return nil
}
switch handlerType {
case "openai":
result := map[string]any{
"id": model.ID,
"object": "model",
"owned_by": model.OwnedBy,
}
if model.Created > 0 {
result["created"] = model.Created
}
if model.Type != "" {
result["type"] = model.Type
}
if model.DisplayName != "" {
result["display_name"] = model.DisplayName
}
if model.Version != "" {
result["version"] = model.Version
}
if model.Description != "" {
result["description"] = model.Description
}
if model.ContextLength > 0 {
result["context_length"] = model.ContextLength
}
if model.MaxContextLength > 0 {
result["max_context_length"] = model.MaxContextLength
}
if model.MaxCompletionTokens > 0 {
result["max_completion_tokens"] = model.MaxCompletionTokens
}
if len(model.SupportedParameters) > 0 {
result["supported_parameters"] = append([]string(nil), model.SupportedParameters...)
}
return result
case "claude":
result := map[string]any{
"id": model.ID,
"object": "model",
"owned_by": model.OwnedBy,
}
if model.Created > 0 {
result["created_at"] = time.Unix(model.Created, 0).UTC().Format(time.RFC3339)
}
result["type"] = "model"
if model.DisplayName != "" {
result["display_name"] = model.DisplayName
} else {
result["display_name"] = model.ID
}
maxInput := model.ContextLength
if maxInput <= 0 {
maxInput = DefaultClaudeMaxInputTokens
}
maxOutput := model.MaxCompletionTokens
if maxOutput <= 0 {
maxOutput = DefaultClaudeMaxOutputTokens
}
result["max_input_tokens"] = maxInput
result["max_tokens"] = maxOutput
return result
case "gemini":
result := map[string]any{}
if model.Name != "" {
result["name"] = model.Name
} else {
result["name"] = model.ID
}
if model.Version != "" {
result["version"] = model.Version
}
if model.DisplayName != "" {
result["displayName"] = model.DisplayName
}
if model.Description != "" {
result["description"] = model.Description
}
inputTokenLimit := model.InputTokenLimit
if inputTokenLimit <= 0 {
inputTokenLimit = model.ContextLength
}
if inputTokenLimit > 0 {
result["inputTokenLimit"] = inputTokenLimit
}
outputTokenLimit := model.OutputTokenLimit
if outputTokenLimit <= 0 {
outputTokenLimit = model.MaxCompletionTokens
}
if outputTokenLimit > 0 {
result["outputTokenLimit"] = outputTokenLimit
}
if len(model.SupportedGenerationMethods) > 0 {
result["supportedGenerationMethods"] = append([]string(nil), model.SupportedGenerationMethods...)
}
if len(model.SupportedInputModalities) > 0 {
result["supportedInputModalities"] = append([]string(nil), model.SupportedInputModalities...)
}
if len(model.SupportedOutputModalities) > 0 {
result["supportedOutputModalities"] = append([]string(nil), model.SupportedOutputModalities...)
}
return result
default:
// Generic format
result := map[string]any{
"id": model.ID,
"object": "model",
}
if model.OwnedBy != "" {
result["owned_by"] = model.OwnedBy
}
if model.Type != "" {
result["type"] = model.Type
}
if model.Created != 0 {
result["created"] = model.Created
}
return result
}
}
// CleanupExpiredQuotas removes expired quota and bounded suspension entries
func (r *ModelRegistry) CleanupExpiredQuotas() {
r.mutex.Lock()
defer r.mutex.Unlock()
now := r.now()
invalidated := false
for modelID, registration := range r.models {
if registration == nil {
continue
}
for clientID, quotaTime := range registration.QuotaExceededClients {
if recoveryAt, ok := clientQuotaRecoveryAt(registration, clientID, quotaTime); ok && !now.Before(recoveryAt) {
delete(registration.QuotaExceededClients, clientID)
invalidated = true
log.Debugf("Cleaned up expired quota tracking for model %s, client %s", modelID, clientID)
}
}
for clientID, deadline := range registration.SuspendedClientDeadlines {
if !deadline.IsZero() && !now.Before(deadline) {
delete(registration.SuspendedClientDeadlines, clientID)
delete(registration.SuspendedClients, clientID)
invalidated = true
log.Debugf("Cleaned up expired suspension for model %s, client %s", modelID, clientID)
}
}
}
if invalidated {
r.invalidateAvailableModelsCacheLocked()
}
}
// GetFirstAvailableModel returns the first available model for the given handler type.
// It prioritizes models by their creation timestamp (newest first) and checks if they have
// available clients that are not suspended or over quota.
//
// Parameters:
// - handlerType: The API handler type (e.g., "openai", "claude", "gemini")
//
// Returns:
// - string: The model ID of the first available model, or empty string if none available
// - error: An error if no models are available
func (r *ModelRegistry) GetFirstAvailableModel(handlerType string) (string, error) {
// Get all available models for this handler type
models := r.GetAvailableModels(handlerType)
if len(models) == 0 {
return "", fmt.Errorf("no models available for handler type: %s", handlerType)
}
// Sort models by creation timestamp (newest first)
sort.Slice(models, func(i, j int) bool {
// Extract created timestamps from map
createdI, okI := models[i]["created"].(int64)
createdJ, okJ := models[j]["created"].(int64)
if !okI || !okJ {
return false
}
return createdI > createdJ
})
// Find the first model with available clients
for _, model := range models {
if modelID, ok := model["id"].(string); ok {
if count := r.GetModelCount(modelID); count > 0 {
return modelID, nil
}
}
}
return "", fmt.Errorf("no available clients for any model in handler type: %s", handlerType)
}
// GetModelsAndEpochForClient atomically returns the models registered for clientID along with the client's current registration epoch.
func (r *ModelRegistry) GetModelsAndEpochForClient(clientID string) ([]*ModelInfo, uint64) {
r.mutex.RLock()
defer r.mutex.RUnlock()
var epoch uint64
if r.clientEpochs != nil {
epoch = r.clientEpochs[clientID]
}
modelIDs, exists := r.clientModels[clientID]
if !exists || len(modelIDs) == 0 {
return nil, epoch
}
// Try to use client-specific model infos first
clientInfos := r.clientModelInfos[clientID]
seen := make(map[string]struct{})
result := make([]*ModelInfo, 0, len(modelIDs))
for _, modelID := range modelIDs {
if _, dup := seen[modelID]; dup {
continue
}
seen[modelID] = struct{}{}
// Prefer client's own model info to preserve original type/owned_by
if clientInfos != nil {
if info, okInfo := clientInfos[modelID]; okInfo && info != nil {
result = append(result, cloneModelInfo(info))
continue
}
}
// Fallback to global registry (for backwards compatibility)
if reg, okReg := r.models[modelID]; okReg && reg.Info != nil {
result = append(result, cloneModelInfo(reg.Info))
}
}
return result, epoch
}
// ClientRegistrationEpoch returns the current registration epoch for clientID.
func (r *ModelRegistry) ClientRegistrationEpoch(clientID string) uint64 {
r.mutex.RLock()
defer r.mutex.RUnlock()
if r.clientEpochs == nil {
return 0
}
return r.clientEpochs[clientID]
}
// GetModelsForClient returns the models registered for a specific client.
// Parameters:
// - clientID: The client identifier (typically auth file name or auth ID)
//
// Returns:
// - []*ModelInfo: List of models registered for this client, nil if client not found
func (r *ModelRegistry) GetModelsForClient(clientID string) []*ModelInfo {
models, _ := r.GetModelsAndEpochForClient(clientID)
return models
}