6158 lines
175 KiB
Go
6158 lines
175 KiB
Go
package auth
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math/rand/v2"
|
|
"net/http"
|
|
"path/filepath"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
|
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
|
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
|
log "github.com/sirupsen/logrus"
|
|
"github.com/tidwall/sjson"
|
|
)
|
|
|
|
// ProviderExecutor defines the contract required by Manager to execute provider calls.
|
|
type ProviderExecutor interface {
|
|
// Identifier returns the provider key handled by this executor.
|
|
Identifier() string
|
|
// Execute handles non-streaming execution and returns the provider response payload.
|
|
Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error)
|
|
// ExecuteStream handles streaming execution and returns a StreamResult containing
|
|
// upstream headers and a channel of provider chunks.
|
|
ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error)
|
|
// Refresh attempts to refresh provider credentials and returns the updated auth state.
|
|
Refresh(ctx context.Context, auth *Auth) (*Auth, error)
|
|
// CountTokens returns the token count for the given request.
|
|
CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error)
|
|
// HttpRequest injects provider credentials into the supplied HTTP request and executes it.
|
|
// Callers must close the response body when non-nil.
|
|
HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error)
|
|
}
|
|
|
|
// RequestAuthPreparer lets an executor update missing auth metadata immediately
|
|
// before a request. Manager serializes and persists returned updates.
|
|
type RequestAuthPreparer interface {
|
|
ShouldPrepareRequestAuth(auth *Auth) bool
|
|
PrepareRequestAuth(ctx context.Context, auth *Auth) (*Auth, error)
|
|
}
|
|
|
|
// ExecutionSessionCloser allows executors to release per-session runtime resources.
|
|
type ExecutionSessionCloser interface {
|
|
CloseExecutionSession(sessionID string)
|
|
}
|
|
|
|
const (
|
|
homeAuthCountMetadataKey = "__cliproxy_home_auth_count"
|
|
// CloseAllExecutionSessionsID asks an executor to release all active execution sessions.
|
|
// Executors that do not support this marker may ignore it.
|
|
CloseAllExecutionSessionsID = "__all_execution_sessions__"
|
|
)
|
|
|
|
// RefreshEvaluator allows runtime state to override refresh decisions.
|
|
type RefreshEvaluator interface {
|
|
ShouldRefresh(now time.Time, auth *Auth) bool
|
|
}
|
|
|
|
const (
|
|
refreshCheckInterval = 5 * time.Second
|
|
refreshMaxConcurrency = 16
|
|
refreshPendingBackoff = time.Minute
|
|
refreshFailureBackoff = 5 * time.Minute
|
|
// refreshIneffectiveBackoff throttles refresh attempts when an executor returns
|
|
// success but the auth still evaluates as needing refresh (e.g. token expiry
|
|
// wasn't updated). Without this guard, the auto-refresh loop can tight-loop and
|
|
// burn CPU at idle.
|
|
refreshIneffectiveBackoff = 30 * time.Second
|
|
quotaBackoffBase = time.Second
|
|
quotaBackoffMax = 30 * time.Minute
|
|
transientErrorCooldown = time.Minute
|
|
)
|
|
|
|
var quotaCooldownDisabled atomic.Bool
|
|
var transientErrorCooldownSeconds atomic.Int64
|
|
|
|
// SetQuotaCooldownDisabled toggles quota cooldown scheduling globally.
|
|
func SetQuotaCooldownDisabled(disable bool) {
|
|
quotaCooldownDisabled.Store(disable)
|
|
}
|
|
|
|
// SetTransientErrorCooldownSeconds configures cooldowns for 408/500/502/503/504.
|
|
// 0 keeps the legacy default; negative values disable transient error cooldowns.
|
|
func SetTransientErrorCooldownSeconds(seconds int) {
|
|
transientErrorCooldownSeconds.Store(int64(seconds))
|
|
}
|
|
|
|
func quotaCooldownDisabledForAuth(auth *Auth) bool {
|
|
return quotaCooldownDisabledForAuthWithConfig(auth, nil)
|
|
}
|
|
|
|
func quotaCooldownDisabledForAuthWithConfig(auth *Auth, cfg *internalconfig.Config) bool {
|
|
if auth != nil {
|
|
if override, ok := auth.DisableCoolingOverride(); ok {
|
|
return override
|
|
}
|
|
if providerCoolingDisabledForAuth(auth, cfg) {
|
|
return true
|
|
}
|
|
}
|
|
if cfg != nil && cfg.DisableCooling {
|
|
return true
|
|
}
|
|
return quotaCooldownDisabled.Load()
|
|
}
|
|
|
|
func providerCoolingDisabledForAuth(auth *Auth, cfg *internalconfig.Config) bool {
|
|
if auth == nil || cfg == nil {
|
|
return false
|
|
}
|
|
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
|
if provider == "" {
|
|
return false
|
|
}
|
|
providerKey := ""
|
|
compatName := ""
|
|
if auth.Attributes != nil {
|
|
providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
|
|
compatName = strings.TrimSpace(auth.Attributes["compat_name"])
|
|
}
|
|
if providerKey == "" && compatName == "" && provider != "openai-compatibility" {
|
|
return false
|
|
}
|
|
if providerKey == "" {
|
|
providerKey = provider
|
|
}
|
|
entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, provider)
|
|
return entry != nil && entry.DisableCooling
|
|
}
|
|
|
|
func nextTransientErrorRetryAfter(now time.Time) time.Time {
|
|
seconds := transientErrorCooldownSeconds.Load()
|
|
if seconds < 0 {
|
|
return time.Time{}
|
|
}
|
|
if seconds == 0 {
|
|
return now.Add(transientErrorCooldown)
|
|
}
|
|
return now.Add(time.Duration(seconds) * time.Second)
|
|
}
|
|
|
|
// Result captures execution outcome used to adjust auth state.
|
|
type Result struct {
|
|
// AuthID references the auth that produced this result.
|
|
AuthID string
|
|
// Provider is copied for convenience when emitting hooks.
|
|
Provider string
|
|
// Model is the upstream model identifier used for the request.
|
|
Model string
|
|
// Success marks whether the execution succeeded.
|
|
Success bool
|
|
// RetryAfter carries a provider supplied retry hint (e.g. 429 retryDelay).
|
|
RetryAfter *time.Duration
|
|
// Error describes the failure when Success is false.
|
|
Error *Error
|
|
}
|
|
|
|
// Selector chooses an auth candidate for execution.
|
|
type Selector interface {
|
|
Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error)
|
|
}
|
|
|
|
type PluginScheduler interface {
|
|
PickAuth(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error)
|
|
}
|
|
|
|
type pluginSchedulerState interface {
|
|
HasScheduler() bool
|
|
}
|
|
|
|
// StoppableSelector is an optional interface for selectors that hold resources.
|
|
// Selectors that implement this interface will have Stop called during shutdown.
|
|
type StoppableSelector interface {
|
|
Selector
|
|
Stop()
|
|
}
|
|
|
|
// Hook captures lifecycle callbacks for observing auth changes.
|
|
type Hook interface {
|
|
// OnAuthRegistered fires when a new auth is registered.
|
|
OnAuthRegistered(ctx context.Context, auth *Auth)
|
|
// OnAuthUpdated fires when an existing auth changes state.
|
|
OnAuthUpdated(ctx context.Context, auth *Auth)
|
|
// OnResult fires when execution result is recorded.
|
|
OnResult(ctx context.Context, result Result)
|
|
}
|
|
|
|
// NoopHook provides optional hook defaults.
|
|
type NoopHook struct{}
|
|
|
|
// OnAuthRegistered implements Hook.
|
|
func (NoopHook) OnAuthRegistered(context.Context, *Auth) {}
|
|
|
|
// OnAuthUpdated implements Hook.
|
|
func (NoopHook) OnAuthUpdated(context.Context, *Auth) {}
|
|
|
|
// OnResult implements Hook.
|
|
func (NoopHook) OnResult(context.Context, Result) {}
|
|
|
|
// Manager orchestrates auth lifecycle, selection, execution, and persistence.
|
|
type Manager struct {
|
|
store Store
|
|
cooldownStore CooldownStateStore
|
|
executors map[string]ProviderExecutor
|
|
selector Selector
|
|
hook Hook
|
|
mu sync.RWMutex
|
|
auths map[string]*Auth
|
|
scheduler *authScheduler
|
|
// pluginScheduler runs outside m.mu before falling back to native selection.
|
|
pluginScheduler PluginScheduler
|
|
// homeRuntimeAuths caches auths returned by Home so websocket sessions can
|
|
// reuse an established upstream credential without dispatching every turn.
|
|
homeRuntimeAuths map[string]map[string]*Auth
|
|
// providerOffsets tracks per-model provider rotation state for multi-provider routing.
|
|
providerOffsets map[string]int
|
|
|
|
// Retry controls request retry behavior.
|
|
requestRetry atomic.Int32
|
|
maxRetryCredentials atomic.Int32
|
|
maxRetryInterval atomic.Int64
|
|
|
|
// oauthModelAlias stores global OAuth model alias mappings (alias -> upstream name) keyed by channel.
|
|
oauthModelAlias atomic.Value
|
|
|
|
// apiKeyModelAlias caches resolved model alias mappings for API-key auths.
|
|
// Keyed by auth.ID, value is alias(lower) -> upstream model (including suffix).
|
|
apiKeyModelAlias atomic.Value
|
|
|
|
// modelPoolOffsets tracks per-auth alias pool rotation state.
|
|
modelPoolOffsets map[string]int
|
|
|
|
// runtimeConfig stores the latest application config for request-time decisions.
|
|
// It is initialized in NewManager; never Load() before first Store().
|
|
runtimeConfig atomic.Value
|
|
|
|
// Optional HTTP RoundTripper provider injected by host.
|
|
rtProvider RoundTripperProvider
|
|
|
|
// Auto refresh state
|
|
refreshCancel context.CancelFunc
|
|
refreshLoop *authAutoRefreshLoop
|
|
|
|
requestPrepareLocks sync.Map
|
|
// refreshLocks serializes credential refresh per auth ID so concurrent
|
|
// 401 recoveries and auto-refresh workers do not race the same refresh_token.
|
|
refreshLocks sync.Map
|
|
}
|
|
|
|
// NewManager constructs a manager with optional custom selector and hook.
|
|
func NewManager(store Store, selector Selector, hook Hook) *Manager {
|
|
if selector == nil {
|
|
selector = &RoundRobinSelector{}
|
|
}
|
|
if hook == nil {
|
|
hook = NoopHook{}
|
|
}
|
|
manager := &Manager{
|
|
store: store,
|
|
executors: make(map[string]ProviderExecutor),
|
|
selector: selector,
|
|
hook: hook,
|
|
auths: make(map[string]*Auth),
|
|
homeRuntimeAuths: make(map[string]map[string]*Auth),
|
|
providerOffsets: make(map[string]int),
|
|
modelPoolOffsets: make(map[string]int),
|
|
}
|
|
// atomic.Value requires non-nil initial value.
|
|
manager.runtimeConfig.Store(&internalconfig.Config{})
|
|
manager.apiKeyModelAlias.Store(apiKeyModelAliasTable(nil))
|
|
manager.scheduler = newAuthScheduler(selector)
|
|
return manager
|
|
}
|
|
|
|
func (m *Manager) SetPluginScheduler(scheduler PluginScheduler) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
m.mu.Lock()
|
|
m.pluginScheduler = scheduler
|
|
m.mu.Unlock()
|
|
}
|
|
|
|
func (m *Manager) hasPluginScheduler() bool {
|
|
if m == nil {
|
|
return false
|
|
}
|
|
m.mu.RLock()
|
|
scheduler := m.pluginScheduler
|
|
m.mu.RUnlock()
|
|
if scheduler == nil {
|
|
return false
|
|
}
|
|
if state, ok := scheduler.(pluginSchedulerState); ok {
|
|
return state.HasScheduler()
|
|
}
|
|
return true
|
|
}
|
|
|
|
func isBuiltInSelector(selector Selector) bool {
|
|
switch selector.(type) {
|
|
case *RoundRobinSelector, *FillFirstSelector:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (m *Manager) syncSchedulerFromSnapshot(auths []*Auth) {
|
|
if m == nil || m.scheduler == nil {
|
|
return
|
|
}
|
|
m.scheduler.rebuild(auths)
|
|
}
|
|
|
|
func (m *Manager) syncScheduler() {
|
|
if m == nil || m.scheduler == nil {
|
|
return
|
|
}
|
|
m.syncSchedulerFromSnapshot(m.snapshotAuths())
|
|
}
|
|
|
|
func (m *Manager) snapshotAuths() []*Auth {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
out := make([]*Auth, 0, len(m.auths))
|
|
for _, a := range m.auths {
|
|
out = append(out, a.Clone())
|
|
}
|
|
return out
|
|
}
|
|
|
|
// RefreshSchedulerEntry re-upserts a single auth into the scheduler so that its
|
|
// supportedModelSet is rebuilt from the current global model registry state.
|
|
// This must be called after models have been registered for a newly added auth,
|
|
// because the initial scheduler.upsertAuth during Register/Update runs before
|
|
// registerModelsForAuth and therefore snapshots an empty model set.
|
|
func (m *Manager) RefreshSchedulerEntry(authID string) {
|
|
if m == nil || m.scheduler == nil || authID == "" {
|
|
return
|
|
}
|
|
m.mu.RLock()
|
|
auth, ok := m.auths[authID]
|
|
if !ok || auth == nil {
|
|
m.mu.RUnlock()
|
|
return
|
|
}
|
|
snapshot := auth.Clone()
|
|
m.mu.RUnlock()
|
|
m.scheduler.upsertAuth(snapshot)
|
|
}
|
|
|
|
// RefreshSchedulerAll rebuilds scheduler entries for every known auth.
|
|
func (m *Manager) RefreshSchedulerAll() {
|
|
if m == nil {
|
|
return
|
|
}
|
|
m.mu.RLock()
|
|
ids := make([]string, 0, len(m.auths))
|
|
for id := range m.auths {
|
|
ids = append(ids, id)
|
|
}
|
|
m.mu.RUnlock()
|
|
for _, id := range ids {
|
|
m.RefreshSchedulerEntry(id)
|
|
}
|
|
}
|
|
|
|
// ReconcileRegistryModelStates aligns per-model runtime state with the current
|
|
// registry snapshot for one auth.
|
|
//
|
|
// Supported models are reset to a clean state because re-registration already
|
|
// cleared the registry-side cooldown/suspension snapshot. ModelStates for
|
|
// models that are no longer present in the registry are pruned entirely so
|
|
// renamed/removed models cannot keep auth-level status stale.
|
|
func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID string) {
|
|
if m == nil || authID == "" {
|
|
return
|
|
}
|
|
|
|
supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID)
|
|
supported := make(map[string]struct{}, len(supportedModels))
|
|
for _, model := range supportedModels {
|
|
if model == nil {
|
|
continue
|
|
}
|
|
modelKey := canonicalModelKey(model.ID)
|
|
if modelKey == "" {
|
|
continue
|
|
}
|
|
supported[modelKey] = struct{}{}
|
|
}
|
|
|
|
var snapshot *Auth
|
|
now := time.Now()
|
|
|
|
m.mu.Lock()
|
|
auth, ok := m.auths[authID]
|
|
if ok && auth != nil && len(auth.ModelStates) > 0 {
|
|
changed := false
|
|
for modelKey, state := range auth.ModelStates {
|
|
baseModel := canonicalModelKey(modelKey)
|
|
if baseModel == "" {
|
|
baseModel = strings.TrimSpace(modelKey)
|
|
}
|
|
if _, supportedModel := supported[baseModel]; !supportedModel {
|
|
// Drop state for models that disappeared from the current registry
|
|
// snapshot. Keeping them around leaks stale errors into auth-level
|
|
// status, management output, and websocket fallback checks.
|
|
delete(auth.ModelStates, modelKey)
|
|
changed = true
|
|
continue
|
|
}
|
|
if state == nil {
|
|
continue
|
|
}
|
|
if modelStateIsClean(state) {
|
|
continue
|
|
}
|
|
resetModelState(state, now)
|
|
changed = true
|
|
}
|
|
if len(auth.ModelStates) == 0 {
|
|
auth.ModelStates = nil
|
|
}
|
|
if changed {
|
|
updateAggregatedAvailability(auth, now)
|
|
if !hasModelError(auth, now) {
|
|
auth.LastError = nil
|
|
auth.StatusMessage = ""
|
|
auth.Status = StatusActive
|
|
}
|
|
auth.UpdatedAt = now
|
|
if errPersist := m.persist(ctx, auth); errPersist != nil {
|
|
logEntryWithRequestID(ctx).WithField("auth_id", auth.ID).Warnf("failed to persist auth changes during model state reconciliation: %v", errPersist)
|
|
}
|
|
snapshot = auth.Clone()
|
|
}
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
if m.scheduler != nil && snapshot != nil {
|
|
m.scheduler.upsertAuth(snapshot)
|
|
}
|
|
}
|
|
|
|
func (m *Manager) SetSelector(selector Selector) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
if selector == nil {
|
|
selector = &RoundRobinSelector{}
|
|
}
|
|
m.mu.Lock()
|
|
m.selector = selector
|
|
m.mu.Unlock()
|
|
if m.scheduler != nil {
|
|
m.scheduler.setSelector(selector)
|
|
m.syncScheduler()
|
|
}
|
|
}
|
|
|
|
// SetStore swaps the underlying persistence store.
|
|
func (m *Manager) SetStore(store Store) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.store = store
|
|
}
|
|
|
|
// SetCooldownStateStore swaps the independent runtime cooldown state store.
|
|
func (m *Manager) SetCooldownStateStore(store CooldownStateStore) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.cooldownStore = store
|
|
}
|
|
|
|
// SetRoundTripperProvider register a provider that returns a per-auth RoundTripper.
|
|
func (m *Manager) SetRoundTripperProvider(p RoundTripperProvider) {
|
|
m.mu.Lock()
|
|
m.rtProvider = p
|
|
m.mu.Unlock()
|
|
}
|
|
|
|
// SetConfig updates the runtime config snapshot used by request-time helpers.
|
|
// Callers should provide the latest config on reload so per-credential alias mapping stays in sync.
|
|
func (m *Manager) SetConfig(cfg *internalconfig.Config) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
if cfg == nil {
|
|
cfg = &internalconfig.Config{}
|
|
}
|
|
m.runtimeConfig.Store(cfg)
|
|
clearedCooldowns := m.clearDisabledCooldownStates(cfg)
|
|
if !cfg.Home.Enabled {
|
|
m.clearHomeRuntimeAuths()
|
|
}
|
|
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
|
if clearedCooldowns {
|
|
m.persistCooldownStates(context.Background())
|
|
}
|
|
}
|
|
|
|
func (m *Manager) cooldownDisabledForAuth(auth *Auth) bool {
|
|
if m == nil {
|
|
return quotaCooldownDisabledForAuth(auth)
|
|
}
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
return quotaCooldownDisabledForAuthWithConfig(auth, cfg)
|
|
}
|
|
|
|
func (m *Manager) clearDisabledCooldownStates(cfg *internalconfig.Config) bool {
|
|
if m == nil {
|
|
return false
|
|
}
|
|
now := time.Now()
|
|
snapshots := make([]*Auth, 0)
|
|
m.mu.Lock()
|
|
for _, auth := range m.auths {
|
|
if auth == nil {
|
|
continue
|
|
}
|
|
if !quotaCooldownDisabledForAuthWithConfig(auth, cfg) && !auth.Disabled && auth.Status != StatusDisabled {
|
|
continue
|
|
}
|
|
if clearCooldownStateForAuth(auth, now) {
|
|
snapshots = append(snapshots, auth.Clone())
|
|
}
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
if m.scheduler != nil {
|
|
for _, snapshot := range snapshots {
|
|
m.scheduler.upsertAuth(snapshot)
|
|
}
|
|
}
|
|
return len(snapshots) > 0
|
|
}
|
|
|
|
// RestoreCooldownStates restores unexpired persisted cooldown records into registered auths.
|
|
func (m *Manager) RestoreCooldownStates(ctx context.Context) error {
|
|
if m == nil {
|
|
return nil
|
|
}
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
m.mu.RLock()
|
|
store := m.cooldownStore
|
|
m.mu.RUnlock()
|
|
if store == nil {
|
|
return nil
|
|
}
|
|
records, errLoad := store.Load(ctx)
|
|
if errLoad != nil {
|
|
return errLoad
|
|
}
|
|
if len(records) == 0 {
|
|
return nil
|
|
}
|
|
|
|
now := time.Now()
|
|
authLevelRecords := make([]CooldownStateRecord, 0)
|
|
snapshotsByID := make(map[string]*Auth)
|
|
|
|
m.mu.Lock()
|
|
for _, record := range records {
|
|
if strings.TrimSpace(record.Model) == "" {
|
|
authLevelRecords = append(authLevelRecords, record)
|
|
continue
|
|
}
|
|
if m.restoreCooldownRecordLocked(record, now) {
|
|
if auth := m.auths[strings.TrimSpace(record.AuthID)]; auth != nil {
|
|
snapshotsByID[auth.ID] = auth.Clone()
|
|
}
|
|
}
|
|
}
|
|
for _, record := range authLevelRecords {
|
|
if m.restoreCooldownRecordLocked(record, now) {
|
|
if auth := m.auths[strings.TrimSpace(record.AuthID)]; auth != nil {
|
|
snapshotsByID[auth.ID] = auth.Clone()
|
|
}
|
|
}
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
if m.scheduler != nil {
|
|
for _, snapshot := range snapshotsByID {
|
|
m.scheduler.upsertAuth(snapshot)
|
|
}
|
|
}
|
|
m.persistCooldownStates(ctx)
|
|
return nil
|
|
}
|
|
|
|
func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now time.Time) bool {
|
|
authID := strings.TrimSpace(record.AuthID)
|
|
if authID == "" || record.NextRetryAfter.IsZero() || !record.NextRetryAfter.After(now) {
|
|
return false
|
|
}
|
|
auth := m.auths[authID]
|
|
if auth == nil || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) {
|
|
return false
|
|
}
|
|
updatedAt := record.UpdatedAt
|
|
if updatedAt.IsZero() {
|
|
updatedAt = now
|
|
}
|
|
reason := strings.TrimSpace(record.Reason)
|
|
model := strings.TrimSpace(record.Model)
|
|
quota := record.Quota
|
|
if quota.Exceeded && quota.NextRecoverAt.IsZero() {
|
|
quota.NextRecoverAt = record.NextRetryAfter
|
|
}
|
|
|
|
if model == "" {
|
|
auth.Unavailable = true
|
|
auth.Status = StatusError
|
|
auth.NextRetryAfter = record.NextRetryAfter
|
|
auth.Quota = quota
|
|
auth.UpdatedAt = updatedAt
|
|
if reason != "" {
|
|
auth.StatusMessage = reason
|
|
}
|
|
auth.LastError = cloneError(record.LastError)
|
|
return true
|
|
}
|
|
|
|
state := ensureModelState(auth, model)
|
|
state.Unavailable = true
|
|
state.Status = StatusError
|
|
state.NextRetryAfter = record.NextRetryAfter
|
|
state.Quota = quota
|
|
state.UpdatedAt = updatedAt
|
|
if reason != "" {
|
|
state.StatusMessage = reason
|
|
}
|
|
state.LastError = cloneError(record.LastError)
|
|
updateAggregatedAvailability(auth, now)
|
|
return true
|
|
}
|
|
|
|
func clearCooldownStateForAuth(auth *Auth, now time.Time) bool {
|
|
if auth == nil {
|
|
return false
|
|
}
|
|
changed := false
|
|
if auth.Unavailable || !auth.NextRetryAfter.IsZero() || auth.Quota.Exceeded || !auth.Quota.NextRecoverAt.IsZero() {
|
|
auth.Unavailable = false
|
|
auth.NextRetryAfter = time.Time{}
|
|
auth.Quota = QuotaState{}
|
|
auth.UpdatedAt = now
|
|
changed = true
|
|
}
|
|
for _, state := range auth.ModelStates {
|
|
if state == nil {
|
|
continue
|
|
}
|
|
if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() {
|
|
state.Unavailable = false
|
|
state.NextRetryAfter = time.Time{}
|
|
state.Quota = QuotaState{}
|
|
state.UpdatedAt = now
|
|
changed = true
|
|
}
|
|
}
|
|
if len(auth.ModelStates) > 0 {
|
|
updateAggregatedAvailability(auth, now)
|
|
}
|
|
return changed
|
|
}
|
|
|
|
func dedupeStrings(values []string) []string {
|
|
if len(values) < 2 {
|
|
return values
|
|
}
|
|
seen := make(map[string]struct{}, len(values))
|
|
out := values[:0]
|
|
for _, value := range values {
|
|
value = strings.TrimSpace(value)
|
|
if value == "" {
|
|
continue
|
|
}
|
|
if _, ok := seen[value]; ok {
|
|
continue
|
|
}
|
|
seen[value] = struct{}{}
|
|
out = append(out, value)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// ResetQuota clears quota/cooldown state for an auth and resumes registry routing.
|
|
func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []string, error) {
|
|
if m == nil {
|
|
return nil, nil, nil
|
|
}
|
|
authID = strings.TrimSpace(authID)
|
|
if authID == "" {
|
|
return nil, nil, fmt.Errorf("auth id is required")
|
|
}
|
|
|
|
now := time.Now()
|
|
var snapshot *Auth
|
|
models := make([]string, 0)
|
|
registeredModels := modelsForRegisteredAuth(authID)
|
|
cooldownStateChanged := false
|
|
|
|
m.mu.Lock()
|
|
auth, ok := m.auths[authID]
|
|
if !ok || auth == nil {
|
|
m.mu.Unlock()
|
|
return nil, nil, nil
|
|
}
|
|
|
|
var cooldownRecordsBefore []CooldownStateRecord
|
|
trackCooldownState := m.cooldownStore != nil
|
|
if trackCooldownState {
|
|
cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now)
|
|
}
|
|
|
|
for modelKey, state := range auth.ModelStates {
|
|
if strings.TrimSpace(modelKey) == "" {
|
|
continue
|
|
}
|
|
models = append(models, modelKey)
|
|
if state != nil {
|
|
resetModelState(state, now)
|
|
}
|
|
}
|
|
if clearCooldownStateForAuth(auth, now) {
|
|
if len(models) == 0 {
|
|
models = append(models, registeredModels...)
|
|
}
|
|
} else if len(auth.ModelStates) > 0 {
|
|
updateAggregatedAvailability(auth, now)
|
|
}
|
|
|
|
if len(models) == 0 {
|
|
models = append(models, registeredModels...)
|
|
}
|
|
models = dedupeStrings(models)
|
|
|
|
if !auth.Disabled && auth.Status != StatusDisabled && !hasModelError(auth, now) {
|
|
auth.LastError = nil
|
|
auth.StatusMessage = ""
|
|
auth.Status = StatusActive
|
|
}
|
|
auth.UpdatedAt = now
|
|
if errPersist := m.persist(ctx, auth); errPersist != nil {
|
|
m.mu.Unlock()
|
|
return nil, nil, errPersist
|
|
}
|
|
snapshot = auth.Clone()
|
|
if trackCooldownState {
|
|
cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now)
|
|
cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter)
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
for _, modelKey := range models {
|
|
registry.GetGlobalRegistry().ClearModelQuotaExceeded(authID, modelKey)
|
|
registry.GetGlobalRegistry().ResumeClientModel(authID, modelKey)
|
|
}
|
|
if m.scheduler != nil && snapshot != nil {
|
|
m.scheduler.upsertAuth(snapshot)
|
|
}
|
|
if snapshot != nil && cooldownStateChanged {
|
|
m.persistCooldownStates(ctx)
|
|
}
|
|
return snapshot, models, nil
|
|
}
|
|
|
|
func modelsForRegisteredAuth(authID string) []string {
|
|
supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID)
|
|
models := make([]string, 0, len(supportedModels))
|
|
for _, supportedModel := range supportedModels {
|
|
if supportedModel == nil || strings.TrimSpace(supportedModel.ID) == "" {
|
|
continue
|
|
}
|
|
models = append(models, supportedModel.ID)
|
|
}
|
|
return models
|
|
}
|
|
|
|
func (m *Manager) persistCooldownStates(ctx context.Context) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
records, store := m.cooldownStateSnapshot()
|
|
if store == nil {
|
|
return
|
|
}
|
|
if errSave := store.Save(ctx, records); errSave != nil {
|
|
logEntryWithRequestID(ctx).Warnf("failed to persist cooldown state: %v", errSave)
|
|
}
|
|
}
|
|
|
|
func (m *Manager) cooldownStateSnapshot() ([]CooldownStateRecord, CooldownStateStore) {
|
|
now := time.Now()
|
|
records := make([]CooldownStateRecord, 0)
|
|
|
|
m.mu.RLock()
|
|
store := m.cooldownStore
|
|
if store == nil {
|
|
m.mu.RUnlock()
|
|
return nil, nil
|
|
}
|
|
for _, auth := range m.auths {
|
|
records = append(records, m.cooldownStateRecordsForAuthLocked(auth, now)...)
|
|
}
|
|
m.mu.RUnlock()
|
|
|
|
sort.Slice(records, func(i, j int) bool {
|
|
if records[i].Provider != records[j].Provider {
|
|
return records[i].Provider < records[j].Provider
|
|
}
|
|
if records[i].AuthID != records[j].AuthID {
|
|
return records[i].AuthID < records[j].AuthID
|
|
}
|
|
return records[i].Model < records[j].Model
|
|
})
|
|
return records, store
|
|
}
|
|
|
|
func (m *Manager) cooldownStateRecordsForAuthLocked(auth *Auth, now time.Time) []CooldownStateRecord {
|
|
if auth == nil || auth.ID == "" || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) {
|
|
return nil
|
|
}
|
|
records := make([]CooldownStateRecord, 0, 1+len(auth.ModelStates))
|
|
if record, ok := authCooldownStateRecord(auth, now); ok {
|
|
records = append(records, record)
|
|
}
|
|
for model, state := range auth.ModelStates {
|
|
if record, ok := modelCooldownStateRecord(auth, model, state, now); ok {
|
|
records = append(records, record)
|
|
}
|
|
}
|
|
sort.Slice(records, func(i, j int) bool {
|
|
return records[i].Model < records[j].Model
|
|
})
|
|
return records
|
|
}
|
|
|
|
func cooldownStateRecordsEqual(a, b []CooldownStateRecord) bool {
|
|
if len(a) != len(b) {
|
|
return false
|
|
}
|
|
for i := range a {
|
|
if !cooldownStateRecordEqual(a[i], b[i]) {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func cooldownStateRecordEqual(a, b CooldownStateRecord) bool {
|
|
if a.Provider != b.Provider ||
|
|
a.AuthID != b.AuthID ||
|
|
a.AuthFile != b.AuthFile ||
|
|
a.Model != b.Model ||
|
|
a.Status != b.Status ||
|
|
a.Reason != b.Reason ||
|
|
!a.NextRetryAfter.Equal(b.NextRetryAfter) ||
|
|
!a.UpdatedAt.Equal(b.UpdatedAt) ||
|
|
!cooldownQuotaEqual(a.Quota, b.Quota) {
|
|
return false
|
|
}
|
|
return cooldownErrorEqual(a.LastError, b.LastError)
|
|
}
|
|
|
|
func cooldownQuotaEqual(a, b QuotaState) bool {
|
|
return a.Exceeded == b.Exceeded &&
|
|
a.Reason == b.Reason &&
|
|
a.BackoffLevel == b.BackoffLevel &&
|
|
a.NextRecoverAt.Equal(b.NextRecoverAt)
|
|
}
|
|
|
|
func cooldownErrorEqual(a, b *Error) bool {
|
|
if a == nil || b == nil {
|
|
return a == b
|
|
}
|
|
return a.Code == b.Code &&
|
|
a.Message == b.Message &&
|
|
a.Retryable == b.Retryable &&
|
|
a.HTTPStatus == b.HTTPStatus
|
|
}
|
|
|
|
func authCooldownStateRecord(auth *Auth, now time.Time) (CooldownStateRecord, bool) {
|
|
if auth == nil || !auth.Unavailable || auth.NextRetryAfter.IsZero() || !auth.NextRetryAfter.After(now) {
|
|
return CooldownStateRecord{}, false
|
|
}
|
|
return CooldownStateRecord{
|
|
Provider: strings.TrimSpace(auth.Provider),
|
|
AuthID: auth.ID,
|
|
AuthFile: cooldownAuthFile(auth),
|
|
Status: "cooling",
|
|
NextRetryAfter: auth.NextRetryAfter,
|
|
Reason: cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError),
|
|
Quota: auth.Quota,
|
|
LastError: cloneError(auth.LastError),
|
|
UpdatedAt: auth.UpdatedAt,
|
|
}, true
|
|
}
|
|
|
|
func modelCooldownStateRecord(auth *Auth, model string, state *ModelState, now time.Time) (CooldownStateRecord, bool) {
|
|
model = strings.TrimSpace(model)
|
|
if auth == nil || state == nil || model == "" || !state.Unavailable || state.NextRetryAfter.IsZero() || !state.NextRetryAfter.After(now) {
|
|
return CooldownStateRecord{}, false
|
|
}
|
|
return CooldownStateRecord{
|
|
Provider: strings.TrimSpace(auth.Provider),
|
|
AuthID: auth.ID,
|
|
AuthFile: cooldownAuthFile(auth),
|
|
Model: model,
|
|
Status: "cooling",
|
|
NextRetryAfter: state.NextRetryAfter,
|
|
Reason: cooldownReason(state.StatusMessage, state.Quota, state.LastError),
|
|
Quota: state.Quota,
|
|
LastError: cloneError(state.LastError),
|
|
UpdatedAt: state.UpdatedAt,
|
|
}, true
|
|
}
|
|
|
|
func cooldownReason(statusMessage string, quota QuotaState, lastErr *Error) string {
|
|
if reason := strings.TrimSpace(quota.Reason); reason != "" {
|
|
return reason
|
|
}
|
|
if statusMessage = strings.TrimSpace(statusMessage); statusMessage != "" {
|
|
return statusMessage
|
|
}
|
|
if lastErr != nil {
|
|
if code := strings.TrimSpace(lastErr.Code); code != "" {
|
|
return code
|
|
}
|
|
if message := strings.TrimSpace(lastErr.Message); message != "" {
|
|
return message
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// HomeEnabled reports whether the home control plane integration is enabled in the runtime config.
|
|
func (m *Manager) HomeEnabled() bool {
|
|
if m == nil {
|
|
return false
|
|
}
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
return cfg != nil && cfg.Home.Enabled
|
|
}
|
|
|
|
func (m *Manager) lookupAPIKeyUpstreamModel(authID, requestedModel string) string {
|
|
if m == nil {
|
|
return ""
|
|
}
|
|
authID = strings.TrimSpace(authID)
|
|
if authID == "" {
|
|
return ""
|
|
}
|
|
requestedModel = strings.TrimSpace(requestedModel)
|
|
if requestedModel == "" {
|
|
return ""
|
|
}
|
|
table, _ := m.apiKeyModelAlias.Load().(apiKeyModelAliasTable)
|
|
if table == nil {
|
|
return ""
|
|
}
|
|
byAlias := table[authID]
|
|
if len(byAlias) == 0 {
|
|
return ""
|
|
}
|
|
key := strings.ToLower(thinking.ParseSuffix(requestedModel).ModelName)
|
|
if key == "" {
|
|
key = strings.ToLower(requestedModel)
|
|
}
|
|
resolved := strings.TrimSpace(byAlias[key])
|
|
if resolved == "" {
|
|
return ""
|
|
}
|
|
return preserveRequestedModelSuffix(requestedModel, resolved)
|
|
}
|
|
|
|
func isAPIKeyAuth(auth *Auth) bool {
|
|
if auth == nil {
|
|
return false
|
|
}
|
|
return auth.AuthKind() == AuthKindAPIKey
|
|
}
|
|
|
|
func isOpenAICompatAPIKeyAuth(auth *Auth) bool {
|
|
if !isAPIKeyAuth(auth) {
|
|
return false
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
|
|
return true
|
|
}
|
|
if auth.Attributes == nil {
|
|
return false
|
|
}
|
|
return strings.TrimSpace(auth.Attributes["compat_name"]) != ""
|
|
}
|
|
|
|
func openAICompatProviderKey(auth *Auth) string {
|
|
if auth == nil {
|
|
return ""
|
|
}
|
|
if auth.Attributes != nil {
|
|
if providerKey := strings.TrimSpace(auth.Attributes["provider_key"]); providerKey != "" {
|
|
return util.OpenAICompatibleProviderKey(providerKey)
|
|
}
|
|
if compatName := strings.TrimSpace(auth.Attributes["compat_name"]); compatName != "" {
|
|
return util.OpenAICompatibleProviderKey(compatName)
|
|
}
|
|
}
|
|
return util.OpenAICompatibleProviderKey(auth.Provider)
|
|
}
|
|
|
|
func openAICompatModelPoolKey(auth *Auth, requestedModel string) string {
|
|
base := strings.TrimSpace(thinking.ParseSuffix(requestedModel).ModelName)
|
|
if base == "" {
|
|
base = strings.TrimSpace(requestedModel)
|
|
}
|
|
return strings.ToLower(strings.TrimSpace(auth.ID)) + "|" + openAICompatProviderKey(auth) + "|" + strings.ToLower(base)
|
|
}
|
|
|
|
func (m *Manager) nextModelPoolOffset(key string, size int) int {
|
|
if m == nil || size <= 1 {
|
|
return 0
|
|
}
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
return 0
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.modelPoolOffsets == nil {
|
|
m.modelPoolOffsets = make(map[string]int)
|
|
}
|
|
offset := m.modelPoolOffsets[key]
|
|
if offset >= 2_147_483_640 {
|
|
offset = 0
|
|
}
|
|
m.modelPoolOffsets[key] = offset + 1
|
|
if size <= 0 {
|
|
return 0
|
|
}
|
|
return offset % size
|
|
}
|
|
|
|
func rotateStrings(values []string, offset int) []string {
|
|
if len(values) <= 1 {
|
|
return values
|
|
}
|
|
if offset <= 0 {
|
|
out := make([]string, len(values))
|
|
copy(out, values)
|
|
return out
|
|
}
|
|
offset = offset % len(values)
|
|
out := make([]string, 0, len(values))
|
|
out = append(out, values[offset:]...)
|
|
out = append(out, values[:offset]...)
|
|
return out
|
|
}
|
|
|
|
func (m *Manager) resolveOpenAICompatUpstreamModelPool(auth *Auth, requestedModel string) []string {
|
|
if m == nil || !isOpenAICompatAPIKeyAuth(auth) {
|
|
return nil
|
|
}
|
|
requestedModel = strings.TrimSpace(requestedModel)
|
|
if requestedModel == "" {
|
|
return nil
|
|
}
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
if cfg == nil {
|
|
cfg = &internalconfig.Config{}
|
|
}
|
|
providerKey := ""
|
|
compatName := ""
|
|
if auth.Attributes != nil {
|
|
providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
|
|
compatName = strings.TrimSpace(auth.Attributes["compat_name"])
|
|
}
|
|
entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider)
|
|
if entry == nil {
|
|
return nil
|
|
}
|
|
return resolveModelAliasPoolFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
|
|
}
|
|
|
|
func preserveRequestedModelSuffix(requestedModel, resolved string) string {
|
|
return preserveResolvedModelSuffix(resolved, thinking.ParseSuffix(requestedModel))
|
|
}
|
|
|
|
func (m *Manager) executionModelCandidates(auth *Auth, routeModel string) []string {
|
|
if auth != nil && auth.Attributes != nil {
|
|
if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" {
|
|
return []string{homeModel}
|
|
}
|
|
}
|
|
requestedModel := rewriteModelForAuth(routeModel, auth)
|
|
requestedModel = m.applyOAuthModelAlias(auth, requestedModel)
|
|
if pool := m.resolveOpenAICompatUpstreamModelPool(auth, requestedModel); len(pool) > 0 {
|
|
if len(pool) == 1 {
|
|
return pool
|
|
}
|
|
offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, requestedModel), len(pool))
|
|
return rotateStrings(pool, offset)
|
|
}
|
|
resolved := m.applyAPIKeyModelAlias(auth, requestedModel)
|
|
if strings.TrimSpace(resolved) == "" {
|
|
resolved = requestedModel
|
|
}
|
|
return []string{resolved}
|
|
}
|
|
|
|
func (m *Manager) selectionModelForAuth(auth *Auth, routeModel string) string {
|
|
requestedModel := rewriteModelForAuth(routeModel, auth)
|
|
if strings.TrimSpace(requestedModel) == "" {
|
|
requestedModel = strings.TrimSpace(routeModel)
|
|
}
|
|
resolvedModel := m.applyOAuthModelAlias(auth, requestedModel)
|
|
if strings.TrimSpace(resolvedModel) == "" {
|
|
resolvedModel = requestedModel
|
|
}
|
|
return resolvedModel
|
|
}
|
|
|
|
func (m *Manager) selectionModelKeyForAuth(auth *Auth, routeModel string) string {
|
|
return canonicalModelKey(m.selectionModelForAuth(auth, routeModel))
|
|
}
|
|
|
|
func (m *Manager) stateModelForExecution(auth *Auth, routeModel, upstreamModel string, pooled bool) string {
|
|
if auth != nil && auth.Attributes != nil {
|
|
if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" {
|
|
if resolved := strings.TrimSpace(upstreamModel); resolved != "" {
|
|
return resolved
|
|
}
|
|
return homeModel
|
|
}
|
|
}
|
|
stateModel := executionResultModel(routeModel, upstreamModel, pooled)
|
|
selectionModel := m.selectionModelForAuth(auth, routeModel)
|
|
if canonicalModelKey(selectionModel) == canonicalModelKey(upstreamModel) && strings.TrimSpace(selectionModel) != "" {
|
|
return strings.TrimSpace(upstreamModel)
|
|
}
|
|
return stateModel
|
|
}
|
|
|
|
func executionResultModel(routeModel, upstreamModel string, pooled bool) string {
|
|
if pooled {
|
|
if resolved := strings.TrimSpace(upstreamModel); resolved != "" {
|
|
return resolved
|
|
}
|
|
}
|
|
if requested := strings.TrimSpace(routeModel); requested != "" {
|
|
return requested
|
|
}
|
|
return strings.TrimSpace(upstreamModel)
|
|
}
|
|
|
|
func (m *Manager) filterExecutionModels(auth *Auth, routeModel string, candidates []string, pooled bool) []string {
|
|
if len(candidates) == 0 {
|
|
return nil
|
|
}
|
|
now := time.Now()
|
|
out := make([]string, 0, len(candidates))
|
|
for _, upstreamModel := range candidates {
|
|
stateModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled)
|
|
blocked, _, _ := isAuthBlockedForModel(auth, stateModel, now)
|
|
if blocked {
|
|
continue
|
|
}
|
|
out = append(out, upstreamModel)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (m *Manager) preparedExecutionModels(auth *Auth, routeModel string) ([]string, bool) {
|
|
candidates := m.executionModelCandidates(auth, routeModel)
|
|
pooled := len(candidates) > 1
|
|
return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled
|
|
}
|
|
|
|
func (m *Manager) preparedExecutionModelsWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) {
|
|
candidates, pooled, aliasResult := m.executionModelCandidatesWithAlias(auth, routeModel)
|
|
return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled, aliasResult
|
|
}
|
|
|
|
func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) {
|
|
requestedModel := rewriteModelForAuth(routeModel, auth)
|
|
aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel)
|
|
upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult)
|
|
|
|
var candidates []string
|
|
if auth != nil && auth.Attributes != nil {
|
|
if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" {
|
|
candidates = []string{homeModel}
|
|
}
|
|
}
|
|
if len(candidates) == 0 {
|
|
if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) > 0 {
|
|
if len(pool) == 1 {
|
|
candidates = pool
|
|
} else {
|
|
offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, upstreamModel), len(pool))
|
|
candidates = rotateStrings(pool, offset)
|
|
}
|
|
} else {
|
|
resolved := m.applyAPIKeyModelAlias(auth, upstreamModel)
|
|
if strings.TrimSpace(resolved) == "" {
|
|
resolved = upstreamModel
|
|
}
|
|
candidates = []string{resolved}
|
|
}
|
|
}
|
|
pooled := len(candidates) > 1
|
|
return candidates, pooled, aliasResult
|
|
}
|
|
|
|
func (m *Manager) resolveExecutionAliasResult(auth *Auth, routeModel string) OAuthModelAliasResult {
|
|
requestedModel := rewriteModelForAuth(routeModel, auth)
|
|
return m.resolveExecutionAliasResultForRequested(auth, requestedModel)
|
|
}
|
|
|
|
func (m *Manager) resolveExecutionAliasResultForRequested(auth *Auth, requestedModel string) OAuthModelAliasResult {
|
|
if auth != nil && auth.AuthKind() == AuthKindAPIKey {
|
|
return m.resolveAPIKeyModelAliasWithResult(auth, requestedModel)
|
|
}
|
|
return m.applyOAuthModelAliasWithResult(auth, requestedModel)
|
|
}
|
|
|
|
func executionAliasPoolModel(auth *Auth, requestedModel string, aliasResult OAuthModelAliasResult) string {
|
|
if auth != nil && auth.AuthKind() == AuthKindAPIKey {
|
|
if strings.TrimSpace(requestedModel) != "" {
|
|
return requestedModel
|
|
}
|
|
}
|
|
if strings.TrimSpace(aliasResult.UpstreamModel) != "" {
|
|
return aliasResult.UpstreamModel
|
|
}
|
|
return requestedModel
|
|
}
|
|
|
|
func (m *Manager) resolveAPIKeyModelAliasWithResult(auth *Auth, requestedModel string) OAuthModelAliasResult {
|
|
if m == nil || auth == nil {
|
|
return OAuthModelAliasResult{}
|
|
}
|
|
requestedModel = strings.TrimSpace(requestedModel)
|
|
if requestedModel == "" {
|
|
return OAuthModelAliasResult{}
|
|
}
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
if cfg == nil {
|
|
cfg = &internalconfig.Config{}
|
|
}
|
|
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
|
var models []modelAliasEntry
|
|
switch provider {
|
|
case "gemini":
|
|
if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil {
|
|
models = asModelAliasEntries(entry.Models)
|
|
}
|
|
case "gemini-interactions":
|
|
if entry := resolveInteractionsAPIKeyConfig(cfg, auth); entry != nil {
|
|
models = asModelAliasEntries(entry.Models)
|
|
}
|
|
case "claude":
|
|
if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil {
|
|
models = asModelAliasEntries(entry.Models)
|
|
}
|
|
case "codex":
|
|
if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil {
|
|
models = asModelAliasEntries(entry.Models)
|
|
}
|
|
case "vertex":
|
|
if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil {
|
|
models = asModelAliasEntries(entry.Models)
|
|
}
|
|
default:
|
|
providerKey := ""
|
|
compatName := ""
|
|
if auth.Attributes != nil {
|
|
providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
|
|
compatName = strings.TrimSpace(auth.Attributes["compat_name"])
|
|
}
|
|
if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
|
|
if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil {
|
|
models = asModelAliasEntries(entry.Models)
|
|
}
|
|
}
|
|
}
|
|
if len(models) == 0 {
|
|
return OAuthModelAliasResult{UpstreamModel: requestedModel}
|
|
}
|
|
result := resolveModelAliasResultFromConfigModels(requestedModel, models)
|
|
if strings.TrimSpace(result.UpstreamModel) == "" {
|
|
return OAuthModelAliasResult{UpstreamModel: requestedModel}
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (m *Manager) prepareExecutionModels(auth *Auth, routeModel string) []string {
|
|
models, _ := m.preparedExecutionModels(auth, routeModel)
|
|
return models
|
|
}
|
|
|
|
func rewriteForceMappedResponse(resp *cliproxyexecutor.Response, aliasResult OAuthModelAliasResult) {
|
|
if resp == nil || !aliasResult.ForceMapping || strings.TrimSpace(aliasResult.OriginalAlias) == "" {
|
|
return
|
|
}
|
|
resp.Payload = rewriteModelInResponse(resp.Payload, aliasResult.OriginalAlias)
|
|
}
|
|
|
|
func rewriteForceMappedStreamChunk(rewriter *StreamRewriter, payload []byte) []byte {
|
|
if rewriter == nil || len(payload) == 0 {
|
|
return payload
|
|
}
|
|
rewritten := rewriter.RewriteChunk(payload)
|
|
if len(rewritten) > 0 {
|
|
return rewritten
|
|
}
|
|
if bytes.Contains(payload, []byte("data:")) {
|
|
if lineWise := rewriteSSEPayloadLines(payload, rewriter.options.RewriteModel); len(lineWise) > 0 {
|
|
return lineWise
|
|
}
|
|
}
|
|
if len(rewriter.pendingBuf) > 0 {
|
|
return nil
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func finishForceMappedStreamChunks(rewriter *StreamRewriter) []byte {
|
|
if rewriter == nil {
|
|
return nil
|
|
}
|
|
return rewriter.Finish()
|
|
}
|
|
|
|
func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeModel string, now time.Time) ([]*Auth, error) {
|
|
if len(auths) == 0 {
|
|
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
|
|
}
|
|
|
|
availableByPriority := make(map[int][]*Auth)
|
|
cooldownCount := 0
|
|
var earliest time.Time
|
|
for _, candidate := range auths {
|
|
checkModel := m.selectionModelForAuth(candidate, routeModel)
|
|
blocked, reason, next := isAuthBlockedForModel(candidate, checkModel, now)
|
|
if !blocked {
|
|
priority := authPriority(candidate)
|
|
availableByPriority[priority] = append(availableByPriority[priority], candidate)
|
|
continue
|
|
}
|
|
if reason == blockReasonCooldown {
|
|
cooldownCount++
|
|
if !next.IsZero() && (earliest.IsZero() || next.Before(earliest)) {
|
|
earliest = next
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(availableByPriority) == 0 {
|
|
if cooldownCount == len(auths) && !earliest.IsZero() {
|
|
providerForError := provider
|
|
if providerForError == "mixed" {
|
|
providerForError = ""
|
|
}
|
|
resetIn := earliest.Sub(now)
|
|
if resetIn < 0 {
|
|
resetIn = 0
|
|
}
|
|
return nil, newModelCooldownError(routeModel, providerForError, resetIn)
|
|
}
|
|
return nil, &Error{Code: "auth_unavailable", Message: "no auth available"}
|
|
}
|
|
|
|
bestPriority := 0
|
|
found := false
|
|
for priority := range availableByPriority {
|
|
if !found || priority > bestPriority {
|
|
bestPriority = priority
|
|
found = true
|
|
}
|
|
}
|
|
|
|
available := availableByPriority[bestPriority]
|
|
if len(available) > 1 {
|
|
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
|
|
}
|
|
return available, nil
|
|
}
|
|
|
|
func selectionArgForSelector(selector Selector, routeModel string) string {
|
|
if isBuiltInSelector(selector) {
|
|
return ""
|
|
}
|
|
return routeModel
|
|
}
|
|
|
|
func schedulerAttributeSensitive(key string) bool {
|
|
key = strings.ToLower(strings.TrimSpace(key))
|
|
normalized := strings.NewReplacer("-", "_", ".", "_", " ", "_").Replace(key)
|
|
compact := strings.NewReplacer("_", "", "-", "", ".", "", " ", "").Replace(key)
|
|
for _, fragment := range []string{
|
|
"api_key",
|
|
"apikey",
|
|
"token",
|
|
"secret",
|
|
"cookie",
|
|
"credential",
|
|
"password",
|
|
"storage",
|
|
"authorization",
|
|
"auth_header",
|
|
"proxy_url",
|
|
} {
|
|
if strings.Contains(key, fragment) || strings.Contains(normalized, fragment) || strings.Contains(compact, fragment) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func schedulerSafeAttributes(src map[string]string) map[string]string {
|
|
if len(src) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]string, len(src))
|
|
for key, value := range src {
|
|
if schedulerAttributeSensitive(key) {
|
|
continue
|
|
}
|
|
out[key] = value
|
|
}
|
|
if len(out) == 0 {
|
|
return nil
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneSchedulerAnyMap(src map[string]any) map[string]any {
|
|
if len(src) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]any, len(src))
|
|
for key, value := range src {
|
|
out[key] = value
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneAuthSlice(auths []*Auth) []*Auth {
|
|
if len(auths) == 0 {
|
|
return nil
|
|
}
|
|
out := make([]*Auth, 0, len(auths))
|
|
for _, auth := range auths {
|
|
if auth == nil {
|
|
continue
|
|
}
|
|
out = append(out, auth.Clone())
|
|
}
|
|
return out
|
|
}
|
|
|
|
func schedulerAuthCandidates(auths []*Auth) []pluginapi.SchedulerAuthCandidate {
|
|
if len(auths) == 0 {
|
|
return nil
|
|
}
|
|
out := make([]pluginapi.SchedulerAuthCandidate, 0, len(auths))
|
|
for _, auth := range auths {
|
|
if auth == nil {
|
|
continue
|
|
}
|
|
out = append(out, pluginapi.SchedulerAuthCandidate{
|
|
ID: auth.ID,
|
|
Provider: strings.ToLower(strings.TrimSpace(auth.Provider)),
|
|
Priority: authPriority(auth),
|
|
Status: string(auth.Status),
|
|
Attributes: schedulerSafeAttributes(auth.Attributes),
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
func schedulerProviders(provider string, providers []string) []string {
|
|
out := make([]string, 0, len(providers)+1)
|
|
seen := make(map[string]struct{}, len(providers)+1)
|
|
addProvider := func(value string) {
|
|
value = strings.ToLower(strings.TrimSpace(value))
|
|
if value == "" || value == "mixed" {
|
|
return
|
|
}
|
|
if _, ok := seen[value]; ok {
|
|
return
|
|
}
|
|
seen[value] = struct{}{}
|
|
out = append(out, value)
|
|
}
|
|
addProvider(provider)
|
|
for _, value := range providers {
|
|
addProvider(value)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func schedulerOptions(opts cliproxyexecutor.Options) pluginapi.SchedulerOptions {
|
|
return pluginapi.SchedulerOptions{
|
|
Headers: cloneHTTPHeader(opts.Headers),
|
|
Metadata: cloneSchedulerAnyMap(opts.Metadata),
|
|
}
|
|
}
|
|
|
|
func pickSchedulerAuthByID(candidates []*Auth, authID string) *Auth {
|
|
authID = strings.TrimSpace(authID)
|
|
if authID == "" {
|
|
return nil
|
|
}
|
|
for _, candidate := range candidates {
|
|
if candidate != nil && candidate.ID == authID {
|
|
return candidate
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func builtinSchedulerStrategy(delegate string) (schedulerStrategy, bool) {
|
|
switch strings.TrimSpace(delegate) {
|
|
case pluginapi.SchedulerBuiltinRoundRobin:
|
|
return schedulerStrategyRoundRobin, true
|
|
case pluginapi.SchedulerBuiltinFillFirst:
|
|
return schedulerStrategyFillFirst, true
|
|
default:
|
|
return schedulerStrategyCustom, false
|
|
}
|
|
}
|
|
|
|
func (m *Manager) pickViaBuiltinScheduler(ctx context.Context, strategy schedulerStrategy, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, bool, error) {
|
|
if m == nil || m.scheduler == nil {
|
|
return nil, false, nil
|
|
}
|
|
providerKey := strings.ToLower(strings.TrimSpace(provider))
|
|
disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata)
|
|
for {
|
|
var selected *Auth
|
|
var errPick error
|
|
if providerKey == "mixed" {
|
|
selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy)
|
|
if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) {
|
|
m.syncScheduler()
|
|
selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy)
|
|
}
|
|
} else {
|
|
selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy)
|
|
if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) {
|
|
m.syncScheduler()
|
|
selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy)
|
|
}
|
|
}
|
|
if errPick != nil {
|
|
return nil, true, errPick
|
|
}
|
|
if selected == nil {
|
|
return nil, true, &Error{Code: "auth_not_found", Message: "selector returned no auth"}
|
|
}
|
|
if disallowFreeAuth && isFreeCodexAuth(selected) {
|
|
if tried == nil {
|
|
tried = make(map[string]struct{})
|
|
}
|
|
tried[selected.ID] = struct{}{}
|
|
continue
|
|
}
|
|
return selected, true, nil
|
|
}
|
|
}
|
|
|
|
func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginScheduler, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}, candidates []*Auth) (*Auth, bool, error) {
|
|
if scheduler == nil || len(candidates) == 0 {
|
|
return nil, false, nil
|
|
}
|
|
providerKey := strings.ToLower(strings.TrimSpace(provider))
|
|
requestProvider := providerKey
|
|
if providerKey == "mixed" {
|
|
requestProvider = ""
|
|
}
|
|
req := pluginapi.SchedulerPickRequest{
|
|
Provider: requestProvider,
|
|
Providers: schedulerProviders(providerKey, providers),
|
|
Model: model,
|
|
Stream: opts.Stream,
|
|
Options: schedulerOptions(opts),
|
|
Candidates: schedulerAuthCandidates(candidates),
|
|
}
|
|
resp, handled, errPick := scheduler.PickAuth(ctx, req)
|
|
if errPick != nil {
|
|
return nil, true, errPick
|
|
}
|
|
if !handled || !resp.Handled {
|
|
return nil, false, nil
|
|
}
|
|
if selected := pickSchedulerAuthByID(candidates, resp.AuthID); selected != nil {
|
|
return selected, true, nil
|
|
}
|
|
|
|
strategy, okStrategy := builtinSchedulerStrategy(resp.DelegateBuiltin)
|
|
if !okStrategy {
|
|
return nil, false, nil
|
|
}
|
|
return m.pickViaBuiltinScheduler(ctx, strategy, providerKey, providers, model, opts, tried)
|
|
}
|
|
|
|
func (m *Manager) authSupportsRouteModel(registryRef *registry.ModelRegistry, auth *Auth, routeModel string) bool {
|
|
if registryRef == nil || auth == nil {
|
|
return true
|
|
}
|
|
routeKey := canonicalModelKey(routeModel)
|
|
if routeKey == "" {
|
|
return true
|
|
}
|
|
if registryRef.ClientSupportsModel(auth.ID, routeKey) {
|
|
return true
|
|
}
|
|
selectionKey := m.selectionModelKeyForAuth(auth, routeModel)
|
|
return selectionKey != "" && selectionKey != routeKey && registryRef.ClientSupportsModel(auth.ID, selectionKey)
|
|
}
|
|
|
|
func discardStreamChunks(ch <-chan cliproxyexecutor.StreamChunk) {
|
|
if ch == nil {
|
|
return
|
|
}
|
|
go func() {
|
|
for range ch {
|
|
}
|
|
}()
|
|
}
|
|
|
|
type streamBootstrapError struct {
|
|
cause error
|
|
headers http.Header
|
|
}
|
|
|
|
func cloneHTTPHeader(headers http.Header) http.Header {
|
|
if headers == nil {
|
|
return nil
|
|
}
|
|
return headers.Clone()
|
|
}
|
|
|
|
func newStreamBootstrapError(err error, headers http.Header) error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
return &streamBootstrapError{
|
|
cause: err,
|
|
headers: cloneHTTPHeader(headers),
|
|
}
|
|
}
|
|
|
|
func (e *streamBootstrapError) Error() string {
|
|
if e == nil || e.cause == nil {
|
|
return ""
|
|
}
|
|
return e.cause.Error()
|
|
}
|
|
|
|
func (e *streamBootstrapError) Unwrap() error {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
return e.cause
|
|
}
|
|
|
|
func (e *streamBootstrapError) Headers() http.Header {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
return cloneHTTPHeader(e.headers)
|
|
}
|
|
|
|
func streamErrorResult(headers http.Header, err error) *cliproxyexecutor.StreamResult {
|
|
ch := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
ch <- cliproxyexecutor.StreamChunk{Err: err}
|
|
close(ch)
|
|
return &cliproxyexecutor.StreamResult{
|
|
Headers: cloneHTTPHeader(headers),
|
|
Chunks: ch,
|
|
}
|
|
}
|
|
|
|
func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamChunk) ([]cliproxyexecutor.StreamChunk, bool, error) {
|
|
if ch == nil {
|
|
return nil, true, nil
|
|
}
|
|
buffered := make([]cliproxyexecutor.StreamChunk, 0, 1)
|
|
for {
|
|
var (
|
|
chunk cliproxyexecutor.StreamChunk
|
|
ok bool
|
|
)
|
|
if ctx != nil {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, false, ctx.Err()
|
|
case chunk, ok = <-ch:
|
|
}
|
|
} else {
|
|
chunk, ok = <-ch
|
|
}
|
|
if !ok {
|
|
return buffered, true, nil
|
|
}
|
|
if chunk.Err != nil {
|
|
return nil, false, chunk.Err
|
|
}
|
|
buffered = append(buffered, chunk)
|
|
if len(chunk.Payload) > 0 {
|
|
return buffered, false, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult) *cliproxyexecutor.StreamResult {
|
|
out := make(chan cliproxyexecutor.StreamChunk)
|
|
go func() {
|
|
defer close(out)
|
|
var failed bool
|
|
forward := true
|
|
var rewriter *StreamRewriter
|
|
if aliasResult.ForceMapping && strings.TrimSpace(aliasResult.OriginalAlias) != "" {
|
|
rewriter = NewStreamRewriter(StreamRewriteOptions{RewriteModel: aliasResult.OriginalAlias})
|
|
}
|
|
emit := func(chunk cliproxyexecutor.StreamChunk) bool {
|
|
if chunk.Err != nil && !failed {
|
|
failed = true
|
|
rerr := &Error{Message: chunk.Err.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](chunk.Err); ok && se != nil {
|
|
rerr.HTTPStatus = se.StatusCode()
|
|
}
|
|
m.MarkResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr})
|
|
}
|
|
if !forward {
|
|
return false
|
|
}
|
|
if chunk.Err != nil {
|
|
if ctx == nil {
|
|
out <- chunk
|
|
return true
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
forward = false
|
|
return false
|
|
case out <- chunk:
|
|
return true
|
|
}
|
|
}
|
|
if len(chunk.Payload) == 0 {
|
|
return true
|
|
}
|
|
payload := rewriteForceMappedStreamChunk(rewriter, chunk.Payload)
|
|
if len(payload) == 0 {
|
|
return true
|
|
}
|
|
chunk.Payload = payload
|
|
if ctx == nil {
|
|
out <- chunk
|
|
return true
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
forward = false
|
|
return false
|
|
case out <- chunk:
|
|
return true
|
|
}
|
|
}
|
|
for _, chunk := range buffered {
|
|
if ok := emit(chunk); !ok {
|
|
discardStreamChunks(remaining)
|
|
return
|
|
}
|
|
}
|
|
for chunk := range remaining {
|
|
if ok := emit(chunk); !ok {
|
|
discardStreamChunks(remaining)
|
|
return
|
|
}
|
|
}
|
|
if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 {
|
|
tailChunk := cliproxyexecutor.StreamChunk{Payload: tail}
|
|
if !emit(tailChunk) {
|
|
return
|
|
}
|
|
}
|
|
if !failed {
|
|
m.MarkResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true})
|
|
}
|
|
}()
|
|
return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}
|
|
}
|
|
|
|
func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult) (*cliproxyexecutor.StreamResult, error) {
|
|
if executor == nil {
|
|
return nil, &Error{Code: "executor_not_found", Message: "executor not registered"}
|
|
}
|
|
ctx = contextWithRequestedModelAlias(ctx, opts, routeModel)
|
|
var lastErr error
|
|
didRefreshOnUnauthorized := false
|
|
for idx, execModel := range execModels {
|
|
resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled)
|
|
execReq := req
|
|
execReq.Model = execModel
|
|
if executionModel != "" {
|
|
execReq.Model = executionModel
|
|
}
|
|
execOpts := opts
|
|
execReq, execOpts = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel))
|
|
streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts)
|
|
if errStream != nil {
|
|
if errCtx := ctx.Err(); errCtx != nil {
|
|
return nil, errCtx
|
|
}
|
|
if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, errStream, didRefreshOnUnauthorized); okRefresh {
|
|
auth = refreshed
|
|
didRefreshOnUnauthorized = true
|
|
streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts)
|
|
if errStream != nil {
|
|
if errCtx := ctx.Err(); errCtx != nil {
|
|
return nil, errCtx
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if errStream != nil {
|
|
rerr := &Error{Message: errStream.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errStream); ok && se != nil {
|
|
rerr.HTTPStatus = se.StatusCode()
|
|
}
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}
|
|
result.RetryAfter = retryAfterFromError(errStream)
|
|
m.MarkResult(ctx, result)
|
|
if isRequestInvalidError(errStream) {
|
|
return nil, errStream
|
|
}
|
|
lastErr = errStream
|
|
continue
|
|
}
|
|
|
|
buffered, closed, bootstrapErr := readStreamBootstrap(ctx, streamResult.Chunks)
|
|
if bootstrapErr != nil {
|
|
if errCtx := ctx.Err(); errCtx != nil {
|
|
discardStreamChunks(streamResult.Chunks)
|
|
return nil, errCtx
|
|
}
|
|
if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, bootstrapErr, didRefreshOnUnauthorized); okRefresh {
|
|
discardStreamChunks(streamResult.Chunks)
|
|
auth = refreshed
|
|
didRefreshOnUnauthorized = true
|
|
retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts)
|
|
if retryErr != nil {
|
|
if errCtx := ctx.Err(); errCtx != nil {
|
|
return nil, errCtx
|
|
}
|
|
bootstrapErr = retryErr
|
|
streamResult = &cliproxyexecutor.StreamResult{}
|
|
} else {
|
|
streamResult = retryStream
|
|
buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks)
|
|
}
|
|
}
|
|
}
|
|
if bootstrapErr != nil {
|
|
if isRequestInvalidError(bootstrapErr) {
|
|
rerr := &Error{Message: bootstrapErr.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](bootstrapErr); ok && se != nil {
|
|
rerr.HTTPStatus = se.StatusCode()
|
|
}
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}
|
|
result.RetryAfter = retryAfterFromError(bootstrapErr)
|
|
m.MarkResult(ctx, result)
|
|
discardStreamChunks(streamResult.Chunks)
|
|
return nil, bootstrapErr
|
|
}
|
|
if idx < len(execModels)-1 {
|
|
rerr := &Error{Message: bootstrapErr.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](bootstrapErr); ok && se != nil {
|
|
rerr.HTTPStatus = se.StatusCode()
|
|
}
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}
|
|
result.RetryAfter = retryAfterFromError(bootstrapErr)
|
|
m.MarkResult(ctx, result)
|
|
discardStreamChunks(streamResult.Chunks)
|
|
lastErr = bootstrapErr
|
|
continue
|
|
}
|
|
rerr := &Error{Message: bootstrapErr.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](bootstrapErr); ok && se != nil {
|
|
rerr.HTTPStatus = se.StatusCode()
|
|
}
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}
|
|
result.RetryAfter = retryAfterFromError(bootstrapErr)
|
|
m.MarkResult(ctx, result)
|
|
discardStreamChunks(streamResult.Chunks)
|
|
return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers)
|
|
}
|
|
|
|
if closed && len(buffered) == 0 {
|
|
emptyErr := &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true}
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: emptyErr}
|
|
m.MarkResult(ctx, result)
|
|
if idx < len(execModels)-1 {
|
|
lastErr = emptyErr
|
|
continue
|
|
}
|
|
return nil, newStreamBootstrapError(emptyErr, streamResult.Headers)
|
|
}
|
|
|
|
remaining := streamResult.Chunks
|
|
if closed {
|
|
closedCh := make(chan cliproxyexecutor.StreamChunk)
|
|
close(closedCh)
|
|
remaining = closedCh
|
|
}
|
|
return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, aliasResult), nil
|
|
}
|
|
if lastErr == nil {
|
|
lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"}
|
|
}
|
|
return nil, lastErr
|
|
}
|
|
|
|
func (m *Manager) rebuildAPIKeyModelAliasFromRuntimeConfig() {
|
|
if m == nil {
|
|
return
|
|
}
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
if cfg == nil {
|
|
cfg = &internalconfig.Config{}
|
|
}
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.rebuildAPIKeyModelAliasLocked(cfg)
|
|
}
|
|
|
|
// RefreshAPIKeyModelAlias rebuilds the API-key model alias table from the current runtime config.
|
|
func (m *Manager) RefreshAPIKeyModelAlias() {
|
|
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
|
}
|
|
|
|
func (m *Manager) rebuildAPIKeyModelAliasLocked(cfg *internalconfig.Config) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
if cfg == nil {
|
|
cfg = &internalconfig.Config{}
|
|
}
|
|
|
|
out := make(apiKeyModelAliasTable)
|
|
for _, auth := range m.auths {
|
|
if auth == nil {
|
|
continue
|
|
}
|
|
if strings.TrimSpace(auth.ID) == "" {
|
|
continue
|
|
}
|
|
if auth.AuthKind() != AuthKindAPIKey {
|
|
continue
|
|
}
|
|
|
|
byAlias := make(map[string]string)
|
|
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
|
switch provider {
|
|
case "gemini":
|
|
if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil {
|
|
compileAPIKeyModelAliasForModels(byAlias, entry.Models)
|
|
}
|
|
case "gemini-interactions":
|
|
if entry := resolveInteractionsAPIKeyConfig(cfg, auth); entry != nil {
|
|
compileAPIKeyModelAliasForModels(byAlias, entry.Models)
|
|
}
|
|
case "claude":
|
|
if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil {
|
|
compileAPIKeyModelAliasForModels(byAlias, entry.Models)
|
|
}
|
|
case "codex":
|
|
if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil {
|
|
compileAPIKeyModelAliasForModels(byAlias, entry.Models)
|
|
}
|
|
case "vertex":
|
|
if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil {
|
|
compileAPIKeyModelAliasForModels(byAlias, entry.Models)
|
|
}
|
|
default:
|
|
// OpenAI-compat uses config selection from auth.Attributes.
|
|
providerKey := ""
|
|
compatName := ""
|
|
if auth.Attributes != nil {
|
|
providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
|
|
compatName = strings.TrimSpace(auth.Attributes["compat_name"])
|
|
}
|
|
if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
|
|
if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil {
|
|
compileAPIKeyModelAliasForModels(byAlias, entry.Models)
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(byAlias) > 0 {
|
|
out[auth.ID] = byAlias
|
|
}
|
|
}
|
|
|
|
m.apiKeyModelAlias.Store(out)
|
|
}
|
|
|
|
func compileAPIKeyModelAliasForModels[T interface {
|
|
GetName() string
|
|
GetAlias() string
|
|
}](out map[string]string, models []T) {
|
|
if out == nil {
|
|
return
|
|
}
|
|
for i := range models {
|
|
alias := strings.TrimSpace(models[i].GetAlias())
|
|
name := strings.TrimSpace(models[i].GetName())
|
|
if alias == "" || name == "" {
|
|
continue
|
|
}
|
|
aliasKey := strings.ToLower(thinking.ParseSuffix(alias).ModelName)
|
|
if aliasKey == "" {
|
|
aliasKey = strings.ToLower(alias)
|
|
}
|
|
// Config priority: first alias wins.
|
|
if _, exists := out[aliasKey]; exists {
|
|
continue
|
|
}
|
|
out[aliasKey] = name
|
|
// Also allow direct lookup by upstream name (case-insensitive), so lookups on already-upstream
|
|
// models remain a cheap no-op.
|
|
nameKey := strings.ToLower(thinking.ParseSuffix(name).ModelName)
|
|
if nameKey == "" {
|
|
nameKey = strings.ToLower(name)
|
|
}
|
|
if nameKey != "" {
|
|
if _, exists := out[nameKey]; !exists {
|
|
out[nameKey] = name
|
|
}
|
|
}
|
|
// Preserve config suffix priority by seeding a base-name lookup when name already has suffix.
|
|
nameResult := thinking.ParseSuffix(name)
|
|
if nameResult.HasSuffix {
|
|
baseKey := strings.ToLower(strings.TrimSpace(nameResult.ModelName))
|
|
if baseKey != "" {
|
|
if _, exists := out[baseKey]; !exists {
|
|
out[baseKey] = name
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// SetRetryConfig updates retry attempts, credential retry limit and cooldown wait interval.
|
|
func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration, maxRetryCredentials int) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
if retry < 0 {
|
|
retry = 0
|
|
}
|
|
if maxRetryCredentials < 0 {
|
|
maxRetryCredentials = 0
|
|
}
|
|
if maxRetryInterval < 0 {
|
|
maxRetryInterval = 0
|
|
}
|
|
m.requestRetry.Store(int32(retry))
|
|
m.maxRetryCredentials.Store(int32(maxRetryCredentials))
|
|
m.maxRetryInterval.Store(maxRetryInterval.Nanoseconds())
|
|
}
|
|
|
|
// RegisterExecutor registers a provider executor with the manager.
|
|
func (m *Manager) RegisterExecutor(executor ProviderExecutor) {
|
|
if executor == nil {
|
|
return
|
|
}
|
|
provider := strings.TrimSpace(executor.Identifier())
|
|
if provider == "" {
|
|
return
|
|
}
|
|
|
|
var replaced ProviderExecutor
|
|
m.mu.Lock()
|
|
replaced = m.executors[provider]
|
|
m.executors[provider] = executor
|
|
m.mu.Unlock()
|
|
|
|
if replaced == nil || replaced == executor {
|
|
return
|
|
}
|
|
if closer, ok := replaced.(ExecutionSessionCloser); ok && closer != nil {
|
|
closer.CloseExecutionSession(CloseAllExecutionSessionsID)
|
|
}
|
|
}
|
|
|
|
// UnregisterExecutor removes the executor associated with the provider key.
|
|
func (m *Manager) UnregisterExecutor(provider string) {
|
|
provider = strings.ToLower(strings.TrimSpace(provider))
|
|
if provider == "" {
|
|
return
|
|
}
|
|
m.mu.Lock()
|
|
delete(m.executors, provider)
|
|
m.mu.Unlock()
|
|
}
|
|
|
|
// Register inserts a new auth entry into the manager.
|
|
func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
|
|
if auth == nil {
|
|
return nil, nil
|
|
}
|
|
if auth.ID == "" {
|
|
auth.ID = uuid.NewString()
|
|
}
|
|
now := time.Now()
|
|
clearedCooldown := false
|
|
if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled {
|
|
clearedCooldown = clearCooldownStateForAuth(auth, now)
|
|
}
|
|
auth.EnsureIndex()
|
|
authClone := auth.Clone()
|
|
m.mu.Lock()
|
|
m.auths[auth.ID] = authClone
|
|
m.mu.Unlock()
|
|
if !shouldDeferAPIKeyModelAliasRebuild(ctx) {
|
|
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
|
}
|
|
if m.scheduler != nil {
|
|
m.scheduler.upsertAuth(authClone)
|
|
}
|
|
m.queueRefreshReschedule(auth.ID)
|
|
_ = m.persist(ctx, auth)
|
|
m.hook.OnAuthRegistered(ctx, auth.Clone())
|
|
if clearedCooldown {
|
|
m.persistCooldownStates(ctx)
|
|
}
|
|
return auth.Clone(), nil
|
|
}
|
|
|
|
// Update replaces an existing auth entry and notifies hooks.
|
|
func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
|
|
if auth == nil || auth.ID == "" {
|
|
return nil, nil
|
|
}
|
|
m.mu.Lock()
|
|
existing, ok := m.auths[auth.ID]
|
|
if !ok || existing == nil {
|
|
m.mu.Unlock()
|
|
return nil, nil
|
|
}
|
|
if !auth.indexAssigned && auth.Index == "" {
|
|
auth.Index = existing.Index
|
|
auth.indexAssigned = existing.indexAssigned
|
|
}
|
|
auth.Success = existing.Success
|
|
auth.Failed = existing.Failed
|
|
auth.recentRequests = existing.recentRequests
|
|
if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled {
|
|
if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 {
|
|
auth.ModelStates = existing.ModelStates
|
|
}
|
|
}
|
|
now := time.Now()
|
|
clearedCooldown := false
|
|
if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled {
|
|
clearedCooldown = clearCooldownStateForAuth(auth, now)
|
|
}
|
|
auth.EnsureIndex()
|
|
authClone := auth.Clone()
|
|
m.auths[auth.ID] = authClone
|
|
m.mu.Unlock()
|
|
if !shouldDeferAPIKeyModelAliasRebuild(ctx) {
|
|
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
|
}
|
|
if m.scheduler != nil {
|
|
m.scheduler.upsertAuth(authClone)
|
|
}
|
|
m.queueRefreshReschedule(auth.ID)
|
|
_ = m.persist(ctx, auth)
|
|
m.hook.OnAuthUpdated(ctx, auth.Clone())
|
|
if clearedCooldown {
|
|
m.persistCooldownStates(ctx)
|
|
}
|
|
return auth.Clone(), nil
|
|
}
|
|
|
|
// Remove deletes an auth from runtime state without persisting.
|
|
// Disk and token-store deletion must be handled by the caller.
|
|
func (m *Manager) Remove(ctx context.Context, id string) {
|
|
if m == nil {
|
|
return
|
|
}
|
|
id = strings.TrimSpace(id)
|
|
if id == "" {
|
|
return
|
|
}
|
|
_ = ctx
|
|
|
|
m.mu.Lock()
|
|
existing := m.auths[id]
|
|
if existing == nil {
|
|
m.mu.Unlock()
|
|
return
|
|
}
|
|
provider := strings.TrimSpace(existing.Provider)
|
|
delete(m.auths, id)
|
|
if m.modelPoolOffsets != nil {
|
|
delete(m.modelPoolOffsets, id)
|
|
}
|
|
for sessionID, sessionAuths := range m.homeRuntimeAuths {
|
|
if sessionAuths == nil {
|
|
continue
|
|
}
|
|
delete(sessionAuths, id)
|
|
if len(sessionAuths) == 0 {
|
|
delete(m.homeRuntimeAuths, sessionID)
|
|
}
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
if !shouldDeferAPIKeyModelAliasRebuild(ctx) {
|
|
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
|
|
}
|
|
if m.scheduler != nil {
|
|
m.scheduler.removeAuth(id)
|
|
}
|
|
m.queueRefreshUnschedule(id)
|
|
m.invalidateSessionAffinity(id)
|
|
|
|
if provider != "" {
|
|
if exec, ok := m.Executor(provider); ok && exec != nil {
|
|
if closer, okCloser := exec.(ExecutionSessionCloser); okCloser {
|
|
closer.CloseExecutionSession(CloseAllExecutionSessionsID)
|
|
}
|
|
}
|
|
}
|
|
m.persistCooldownStates(ctx)
|
|
}
|
|
|
|
func (m *Manager) invalidateSessionAffinity(authID string) {
|
|
if m == nil || authID == "" {
|
|
return
|
|
}
|
|
if invalidator, ok := m.selector.(interface{ InvalidateAuth(string) }); ok && invalidator != nil {
|
|
invalidator.InvalidateAuth(authID)
|
|
}
|
|
}
|
|
|
|
// Load resets manager state from the backing store.
|
|
func (m *Manager) Load(ctx context.Context) error {
|
|
m.mu.Lock()
|
|
if m.store == nil {
|
|
m.mu.Unlock()
|
|
return nil
|
|
}
|
|
items, err := m.store.List(ctx)
|
|
if err != nil {
|
|
m.mu.Unlock()
|
|
return err
|
|
}
|
|
m.auths = make(map[string]*Auth, len(items))
|
|
for _, auth := range items {
|
|
if auth == nil || auth.ID == "" {
|
|
continue
|
|
}
|
|
auth.EnsureIndex()
|
|
m.auths[auth.ID] = auth.Clone()
|
|
}
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
if cfg == nil {
|
|
cfg = &internalconfig.Config{}
|
|
}
|
|
m.rebuildAPIKeyModelAliasLocked(cfg)
|
|
m.mu.Unlock()
|
|
m.syncScheduler()
|
|
return nil
|
|
}
|
|
|
|
// Execute performs a non-streaming execution using the configured selector and executor.
|
|
// It supports multiple providers for the same model and round-robins the starting provider per model.
|
|
func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
normalized := m.normalizeProviders(providers)
|
|
if len(normalized) == 0 {
|
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
|
}
|
|
|
|
_, maxRetryCredentials, maxWait := m.retrySettings()
|
|
|
|
var lastErr error
|
|
retryModel := authSelectionModelFromOptions(opts, req.Model)
|
|
for attempt := 0; ; attempt++ {
|
|
resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials)
|
|
if errExec == nil {
|
|
return resp, nil
|
|
}
|
|
lastErr = errExec
|
|
wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait)
|
|
if !shouldRetry {
|
|
break
|
|
}
|
|
if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil {
|
|
return cliproxyexecutor.Response{}, errWait
|
|
}
|
|
}
|
|
if lastErr != nil {
|
|
if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) {
|
|
if resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil {
|
|
return cliproxyexecutor.Response{}, errCredits
|
|
} else if ok {
|
|
return resp, nil
|
|
}
|
|
}
|
|
return cliproxyexecutor.Response{}, lastErr
|
|
}
|
|
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
|
|
// It supports multiple providers for the same model and round-robins the starting provider per model.
|
|
func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
normalized := m.normalizeProviders(providers)
|
|
if len(normalized) == 0 {
|
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
|
}
|
|
|
|
_, maxRetryCredentials, maxWait := m.retrySettings()
|
|
|
|
var lastErr error
|
|
retryModel := authSelectionModelFromOptions(opts, req.Model)
|
|
for attempt := 0; ; attempt++ {
|
|
resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials)
|
|
if errExec == nil {
|
|
return resp, nil
|
|
}
|
|
lastErr = errExec
|
|
wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait)
|
|
if !shouldRetry {
|
|
break
|
|
}
|
|
if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil {
|
|
return cliproxyexecutor.Response{}, errWait
|
|
}
|
|
}
|
|
if lastErr != nil {
|
|
return cliproxyexecutor.Response{}, lastErr
|
|
}
|
|
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
|
|
// ExecuteStream performs a streaming execution using the configured selector and executor.
|
|
// It supports multiple providers for the same model and round-robins the starting provider per model.
|
|
func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
normalized := m.normalizeProviders(providers)
|
|
if len(normalized) == 0 {
|
|
return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
|
}
|
|
|
|
_, maxRetryCredentials, maxWait := m.retrySettings()
|
|
|
|
var lastErr error
|
|
retryModel := authSelectionModelFromOptions(opts, req.Model)
|
|
for attempt := 0; ; attempt++ {
|
|
result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials)
|
|
if errStream == nil {
|
|
return result, nil
|
|
}
|
|
lastErr = errStream
|
|
wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, retryModel, maxWait)
|
|
if !shouldRetry {
|
|
break
|
|
}
|
|
if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil {
|
|
return nil, errWait
|
|
}
|
|
}
|
|
if lastErr != nil {
|
|
if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) {
|
|
if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil {
|
|
return nil, errCredits
|
|
} else if ok {
|
|
return result, nil
|
|
}
|
|
}
|
|
var bootstrapErr *streamBootstrapError
|
|
if errors.As(lastErr, &bootstrapErr) && bootstrapErr != nil {
|
|
return streamErrorResult(bootstrapErr.Headers(), bootstrapErr.cause), nil
|
|
}
|
|
return nil, lastErr
|
|
}
|
|
return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
|
|
type requestToFormatResolver interface {
|
|
RequestToFormat(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format
|
|
}
|
|
|
|
func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExecutor, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, requestedModel string) (cliproxyexecutor.Request, cliproxyexecutor.Options) {
|
|
if opts.RequestAfterAuthInterceptor == nil {
|
|
return req, opts
|
|
}
|
|
toFormat := requestToFormat(provider, executor, req, opts)
|
|
resp := opts.RequestAfterAuthInterceptor(ctx, cliproxyexecutor.RequestAfterAuthInterceptRequest{
|
|
SourceFormat: opts.SourceFormat,
|
|
ToFormat: toFormat,
|
|
Model: req.Model,
|
|
RequestedModel: requestedModel,
|
|
Stream: opts.Stream,
|
|
Headers: cloneRequestHeaders(opts.Headers),
|
|
Body: bytes.Clone(req.Payload),
|
|
Metadata: opts.Metadata,
|
|
})
|
|
opts.Headers = mergeRequestHeaders(opts.Headers, resp.Headers, resp.ClearHeaders)
|
|
if len(resp.Body) > 0 {
|
|
req.Payload = bytes.Clone(resp.Body)
|
|
opts.OriginalRequest = bytes.Clone(resp.Body)
|
|
}
|
|
return req, opts
|
|
}
|
|
|
|
func requestToFormat(provider string, executor ProviderExecutor, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format {
|
|
resolver, ok := executor.(requestToFormatResolver)
|
|
if ok && resolver != nil {
|
|
formatRequestTo := resolver.RequestToFormat(req, opts)
|
|
if formatRequestTo != "" {
|
|
return formatRequestTo
|
|
}
|
|
}
|
|
source := opts.SourceFormat.String()
|
|
if source == "openai-image" || source == "openai-video" {
|
|
return opts.SourceFormat
|
|
}
|
|
if opts.Alt == "responses/compact" && !opts.Stream {
|
|
return sdktranslator.FormatOpenAIResponse
|
|
}
|
|
switch strings.ToLower(strings.TrimSpace(provider)) {
|
|
case "codex":
|
|
return sdktranslator.FormatCodex
|
|
case "xai":
|
|
return sdktranslator.FormatCodex
|
|
case "claude":
|
|
return sdktranslator.FormatClaude
|
|
case "gemini", "vertex", "aistudio":
|
|
return sdktranslator.FormatGemini
|
|
case "kimi":
|
|
return sdktranslator.FormatOpenAI
|
|
case "antigravity":
|
|
return sdktranslator.FormatAntigravity
|
|
default:
|
|
return sdktranslator.FormatOpenAI
|
|
}
|
|
}
|
|
|
|
func cloneRequestHeaders(src http.Header) http.Header {
|
|
if src == nil {
|
|
return nil
|
|
}
|
|
dst := make(http.Header, len(src))
|
|
for key, values := range src {
|
|
dst[key] = append([]string(nil), values...)
|
|
}
|
|
return dst
|
|
}
|
|
|
|
func mergeRequestHeaders(current, updates http.Header, clear []string) http.Header {
|
|
if updates == nil && len(clear) == 0 {
|
|
return current
|
|
}
|
|
out := cloneRequestHeaders(current)
|
|
if out == nil && (len(updates) > 0 || len(clear) > 0) {
|
|
out = make(http.Header)
|
|
}
|
|
for _, key := range clear {
|
|
out.Del(key)
|
|
}
|
|
for key, values := range updates {
|
|
out.Del(key)
|
|
for _, value := range values {
|
|
out.Add(key, value)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) {
|
|
if len(providers) == 0 {
|
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
|
}
|
|
routeModel := authSelectionModelFromOptions(opts, req.Model)
|
|
executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model)
|
|
opts = ensureRequestedModelMetadata(opts, routeModel)
|
|
homeMode := m.HomeEnabled()
|
|
homeAuthCount := 1
|
|
tried := make(map[string]struct{})
|
|
attempted := make(map[string]struct{})
|
|
var lastErr error
|
|
for {
|
|
if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
|
|
if lastErr != nil {
|
|
return cliproxyexecutor.Response{}, lastErr
|
|
}
|
|
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
pickOpts := opts
|
|
if homeMode {
|
|
pickOpts = withHomeAuthCount(opts, homeAuthCount)
|
|
}
|
|
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
|
|
if errPick != nil {
|
|
if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) {
|
|
return cliproxyexecutor.Response{}, lastErr
|
|
}
|
|
return cliproxyexecutor.Response{}, errPick
|
|
}
|
|
|
|
entry := logEntryWithRequestID(ctx)
|
|
debugLogAuthSelection(entry, auth, provider, routeModel)
|
|
publishSelectedAuthMetadata(opts.Metadata, auth.ID)
|
|
|
|
tried[auth.ID] = struct{}{}
|
|
execCtx := ctx
|
|
if rt := m.roundTripperFor(auth); rt != nil {
|
|
execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
|
|
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
|
|
}
|
|
execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel)
|
|
|
|
models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel)
|
|
if len(models) == 0 {
|
|
continue
|
|
}
|
|
attempted[auth.ID] = struct{}{}
|
|
var errPrepare error
|
|
auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)
|
|
if errPrepare != nil {
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: &Error{Message: errPrepare.Error()}}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errPrepare); ok && se != nil {
|
|
result.Error.HTTPStatus = se.StatusCode()
|
|
}
|
|
m.MarkResult(execCtx, result)
|
|
lastErr = errPrepare
|
|
continue
|
|
}
|
|
var authErr error
|
|
didRefreshOnUnauthorized := false
|
|
for _, upstreamModel := range models {
|
|
resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled)
|
|
execReq := req
|
|
execReq.Model = upstreamModel
|
|
if restoreExecutionModel {
|
|
execReq.Model = executionModel
|
|
}
|
|
execOpts := opts
|
|
execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel))
|
|
resp, errExec := executor.Execute(execCtx, auth, execReq, execOpts)
|
|
if errExec != nil {
|
|
if errCtx := execCtx.Err(); errCtx != nil {
|
|
return cliproxyexecutor.Response{}, errCtx
|
|
}
|
|
if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh {
|
|
auth = refreshed
|
|
didRefreshOnUnauthorized = true
|
|
resp, errExec = executor.Execute(execCtx, auth, execReq, execOpts)
|
|
if errExec != nil {
|
|
if errCtx := execCtx.Err(); errCtx != nil {
|
|
return cliproxyexecutor.Response{}, errCtx
|
|
}
|
|
}
|
|
}
|
|
}
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil}
|
|
if errExec != nil {
|
|
result.Error = &Error{Message: errExec.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errExec); ok && se != nil {
|
|
result.Error.HTTPStatus = se.StatusCode()
|
|
}
|
|
if ra := retryAfterFromError(errExec); ra != nil {
|
|
result.RetryAfter = ra
|
|
}
|
|
m.MarkResult(execCtx, result)
|
|
if isRequestInvalidError(errExec) {
|
|
return cliproxyexecutor.Response{}, errExec
|
|
}
|
|
authErr = errExec
|
|
continue
|
|
}
|
|
m.MarkResult(execCtx, result)
|
|
rewriteForceMappedResponse(&resp, aliasResult)
|
|
return resp, nil
|
|
}
|
|
if authErr != nil {
|
|
if isRequestInvalidError(authErr) {
|
|
return cliproxyexecutor.Response{}, authErr
|
|
}
|
|
lastErr = authErr
|
|
if homeMode {
|
|
homeAuthCount++
|
|
}
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) {
|
|
if len(providers) == 0 {
|
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
|
}
|
|
routeModel := authSelectionModelFromOptions(opts, req.Model)
|
|
executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model)
|
|
opts = ensureRequestedModelMetadata(opts, routeModel)
|
|
homeMode := m.HomeEnabled()
|
|
homeAuthCount := 1
|
|
tried := make(map[string]struct{})
|
|
attempted := make(map[string]struct{})
|
|
var lastErr error
|
|
for {
|
|
if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
|
|
if lastErr != nil {
|
|
return cliproxyexecutor.Response{}, lastErr
|
|
}
|
|
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
pickOpts := opts
|
|
if homeMode {
|
|
pickOpts = withHomeAuthCount(opts, homeAuthCount)
|
|
}
|
|
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
|
|
if errPick != nil {
|
|
if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) {
|
|
return cliproxyexecutor.Response{}, lastErr
|
|
}
|
|
return cliproxyexecutor.Response{}, errPick
|
|
}
|
|
|
|
entry := logEntryWithRequestID(ctx)
|
|
debugLogAuthSelection(entry, auth, provider, routeModel)
|
|
publishSelectedAuthMetadata(opts.Metadata, auth.ID)
|
|
|
|
tried[auth.ID] = struct{}{}
|
|
execCtx := ctx
|
|
if rt := m.roundTripperFor(auth); rt != nil {
|
|
execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
|
|
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
|
|
}
|
|
execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel)
|
|
|
|
models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel)
|
|
if len(models) == 0 {
|
|
continue
|
|
}
|
|
attempted[auth.ID] = struct{}{}
|
|
var errPrepare error
|
|
auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)
|
|
if errPrepare != nil {
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: &Error{Message: errPrepare.Error()}}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errPrepare); ok && se != nil {
|
|
result.Error.HTTPStatus = se.StatusCode()
|
|
}
|
|
m.MarkResult(execCtx, result)
|
|
lastErr = errPrepare
|
|
continue
|
|
}
|
|
var authErr error
|
|
didRefreshOnUnauthorized := false
|
|
for _, upstreamModel := range models {
|
|
resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled)
|
|
execReq := req
|
|
execReq.Model = upstreamModel
|
|
if restoreExecutionModel {
|
|
execReq.Model = executionModel
|
|
}
|
|
execOpts := opts
|
|
execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel))
|
|
resp, errExec := executor.CountTokens(execCtx, auth, execReq, execOpts)
|
|
if errExec != nil {
|
|
if errCtx := execCtx.Err(); errCtx != nil {
|
|
return cliproxyexecutor.Response{}, errCtx
|
|
}
|
|
if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh {
|
|
auth = refreshed
|
|
didRefreshOnUnauthorized = true
|
|
resp, errExec = executor.CountTokens(execCtx, auth, execReq, execOpts)
|
|
if errExec != nil {
|
|
if errCtx := execCtx.Err(); errCtx != nil {
|
|
return cliproxyexecutor.Response{}, errCtx
|
|
}
|
|
}
|
|
}
|
|
}
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil}
|
|
if errExec != nil {
|
|
result.Error = &Error{Message: errExec.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errExec); ok && se != nil {
|
|
result.Error.HTTPStatus = se.StatusCode()
|
|
}
|
|
if ra := retryAfterFromError(errExec); ra != nil {
|
|
result.RetryAfter = ra
|
|
}
|
|
m.MarkResult(execCtx, result)
|
|
if isRequestInvalidError(errExec) {
|
|
return cliproxyexecutor.Response{}, errExec
|
|
}
|
|
authErr = errExec
|
|
continue
|
|
}
|
|
m.MarkResult(execCtx, result)
|
|
rewriteForceMappedResponse(&resp, aliasResult)
|
|
return resp, nil
|
|
}
|
|
if authErr != nil {
|
|
if isRequestInvalidError(authErr) {
|
|
return cliproxyexecutor.Response{}, authErr
|
|
}
|
|
lastErr = authErr
|
|
if homeMode {
|
|
homeAuthCount++
|
|
}
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (*cliproxyexecutor.StreamResult, error) {
|
|
if len(providers) == 0 {
|
|
return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
|
}
|
|
routeModel := authSelectionModelFromOptions(opts, req.Model)
|
|
executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model)
|
|
opts = ensureRequestedModelMetadata(opts, routeModel)
|
|
homeMode := m.HomeEnabled()
|
|
homeAuthCount := 1
|
|
tried := make(map[string]struct{})
|
|
attempted := make(map[string]struct{})
|
|
var lastErr error
|
|
for {
|
|
if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials {
|
|
if lastErr != nil {
|
|
return nil, lastErr
|
|
}
|
|
return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
pickOpts := opts
|
|
if homeMode {
|
|
pickOpts = withHomeAuthCount(opts, homeAuthCount)
|
|
}
|
|
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried)
|
|
if errPick != nil {
|
|
if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) {
|
|
return nil, lastErr
|
|
}
|
|
return nil, errPick
|
|
}
|
|
|
|
entry := logEntryWithRequestID(ctx)
|
|
debugLogAuthSelection(entry, auth, provider, routeModel)
|
|
publishSelectedAuthMetadata(opts.Metadata, auth.ID)
|
|
|
|
tried[auth.ID] = struct{}{}
|
|
execCtx := ctx
|
|
if rt := m.roundTripperFor(auth); rt != nil {
|
|
execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt)
|
|
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
|
|
}
|
|
models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel)
|
|
if len(models) == 0 {
|
|
continue
|
|
}
|
|
attempted[auth.ID] = struct{}{}
|
|
var errPrepare error
|
|
auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)
|
|
if errPrepare != nil {
|
|
result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: &Error{Message: errPrepare.Error()}}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errPrepare); ok && se != nil {
|
|
result.Error.HTTPStatus = se.StatusCode()
|
|
}
|
|
m.MarkResult(execCtx, result)
|
|
lastErr = errPrepare
|
|
continue
|
|
}
|
|
execReq := sanitizeDownstreamWebsocketFallbackRequest(execCtx, auth, req)
|
|
streamExecutionModel := ""
|
|
if restoreExecutionModel {
|
|
streamExecutionModel = executionModel
|
|
}
|
|
streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, opts, routeModel, streamExecutionModel, models, pooled, aliasResult)
|
|
if errStream != nil {
|
|
if errCtx := execCtx.Err(); errCtx != nil {
|
|
return nil, errCtx
|
|
}
|
|
if isRequestInvalidError(errStream) {
|
|
return nil, errStream
|
|
}
|
|
lastErr = errStream
|
|
if homeMode {
|
|
homeAuthCount++
|
|
}
|
|
continue
|
|
}
|
|
return streamResult, nil
|
|
}
|
|
}
|
|
|
|
func sanitizeDownstreamWebsocketFallbackRequest(ctx context.Context, auth *Auth, req cliproxyexecutor.Request) cliproxyexecutor.Request {
|
|
if !cliproxyexecutor.DownstreamWebsocket(ctx) || authWebsocketsEnabled(auth) || len(req.Payload) == 0 {
|
|
return req
|
|
}
|
|
updated, errDelete := sjson.DeleteBytes(req.Payload, "generate")
|
|
if errDelete != nil {
|
|
return req
|
|
}
|
|
req.Payload = updated
|
|
return req
|
|
}
|
|
|
|
func ensureRequestedModelMetadata(opts cliproxyexecutor.Options, requestedModel string) cliproxyexecutor.Options {
|
|
requestedModel = strings.TrimSpace(requestedModel)
|
|
if requestedModel == "" {
|
|
return opts
|
|
}
|
|
if hasRequestedModelMetadata(opts.Metadata) {
|
|
return opts
|
|
}
|
|
if len(opts.Metadata) == 0 {
|
|
opts.Metadata = map[string]any{cliproxyexecutor.RequestedModelMetadataKey: requestedModel}
|
|
return opts
|
|
}
|
|
meta := make(map[string]any, len(opts.Metadata)+1)
|
|
for k, v := range opts.Metadata {
|
|
meta[k] = v
|
|
}
|
|
meta[cliproxyexecutor.RequestedModelMetadataKey] = requestedModel
|
|
opts.Metadata = meta
|
|
return opts
|
|
}
|
|
|
|
func authSelectionModelFromOptions(opts cliproxyexecutor.Options, fallback string) string {
|
|
fallback = strings.TrimSpace(fallback)
|
|
if len(opts.Metadata) == 0 {
|
|
return fallback
|
|
}
|
|
raw, ok := opts.Metadata[cliproxyexecutor.AuthSelectionModelMetadataKey]
|
|
if !ok || raw == nil {
|
|
return fallback
|
|
}
|
|
switch value := raw.(type) {
|
|
case string:
|
|
if strings.TrimSpace(value) != "" {
|
|
return strings.TrimSpace(value)
|
|
}
|
|
case []byte:
|
|
if strings.TrimSpace(string(value)) != "" {
|
|
return strings.TrimSpace(string(value))
|
|
}
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
func executionModelForAuthSelection(opts cliproxyexecutor.Options, model string) (string, bool) {
|
|
model = strings.TrimSpace(model)
|
|
if model == "" {
|
|
return "", false
|
|
}
|
|
selectionModel := authSelectionModelFromOptions(opts, model)
|
|
if selectionModel == model {
|
|
return "", false
|
|
}
|
|
return model, true
|
|
}
|
|
|
|
func withHomeAuthCount(opts cliproxyexecutor.Options, count int) cliproxyexecutor.Options {
|
|
if count <= 0 {
|
|
count = 1
|
|
}
|
|
meta := make(map[string]any, len(opts.Metadata)+1)
|
|
for k, v := range opts.Metadata {
|
|
meta[k] = v
|
|
}
|
|
meta[homeAuthCountMetadataKey] = count
|
|
opts.Metadata = meta
|
|
return opts
|
|
}
|
|
|
|
func homeAuthCountFromMetadata(meta map[string]any) int {
|
|
if len(meta) == 0 {
|
|
return 1
|
|
}
|
|
switch value := meta[homeAuthCountMetadataKey].(type) {
|
|
case int:
|
|
if value > 0 {
|
|
return value
|
|
}
|
|
case int64:
|
|
if value > 0 {
|
|
return int(value)
|
|
}
|
|
case float64:
|
|
if value > 0 {
|
|
return int(value)
|
|
}
|
|
}
|
|
return 1
|
|
}
|
|
|
|
func hasRequestedModelMetadata(meta map[string]any) bool {
|
|
if len(meta) == 0 {
|
|
return false
|
|
}
|
|
raw, ok := meta[cliproxyexecutor.RequestedModelMetadataKey]
|
|
if !ok || raw == nil {
|
|
return false
|
|
}
|
|
switch v := raw.(type) {
|
|
case string:
|
|
return strings.TrimSpace(v) != ""
|
|
case []byte:
|
|
return strings.TrimSpace(string(v)) != ""
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
type requestAuthPrepareLock struct {
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func (m *Manager) prepareRequestAuth(ctx context.Context, executor ProviderExecutor, auth *Auth) (*Auth, error) {
|
|
if m == nil || executor == nil || auth == nil {
|
|
return auth, nil
|
|
}
|
|
preparer, ok := executor.(RequestAuthPreparer)
|
|
if !ok || preparer == nil || !preparer.ShouldPrepareRequestAuth(auth) {
|
|
return auth, nil
|
|
}
|
|
|
|
id := strings.TrimSpace(auth.ID)
|
|
if id == "" {
|
|
return preparer.PrepareRequestAuth(ctx, auth.Clone())
|
|
}
|
|
|
|
lockValue, _ := m.requestPrepareLocks.LoadOrStore(id, &requestAuthPrepareLock{})
|
|
lock, ok := lockValue.(*requestAuthPrepareLock)
|
|
if !ok || lock == nil {
|
|
return preparer.PrepareRequestAuth(ctx, auth.Clone())
|
|
}
|
|
|
|
lock.mu.Lock()
|
|
defer lock.mu.Unlock()
|
|
|
|
target := auth.Clone()
|
|
m.mu.RLock()
|
|
if current := m.auths[id]; current != nil {
|
|
target = current.Clone()
|
|
}
|
|
m.mu.RUnlock()
|
|
|
|
if !preparer.ShouldPrepareRequestAuth(target) {
|
|
return target, nil
|
|
}
|
|
|
|
updated, errPrepare := preparer.PrepareRequestAuth(ctx, target)
|
|
if errPrepare != nil {
|
|
return auth, errPrepare
|
|
}
|
|
if updated == nil {
|
|
return target, nil
|
|
}
|
|
|
|
saved, errUpdate := m.Update(ctx, updated)
|
|
if errUpdate != nil {
|
|
return updated, errUpdate
|
|
}
|
|
if saved != nil {
|
|
return saved, nil
|
|
}
|
|
return updated, nil
|
|
}
|
|
|
|
func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.Options, fallback string) context.Context {
|
|
alias := requestedModelAliasFromOptions(opts, fallback)
|
|
ctx = coreusage.WithRequestedModelAlias(ctx, alias)
|
|
effort := reasoningEffortFromOptions(opts)
|
|
if effort != "" {
|
|
ctx = coreusage.WithReasoningEffort(ctx, effort)
|
|
}
|
|
serviceTier := serviceTierFromOptions(opts)
|
|
if serviceTier != "" {
|
|
ctx = coreusage.WithServiceTier(ctx, serviceTier)
|
|
}
|
|
return ctx
|
|
}
|
|
|
|
func requestedModelAliasFromOptions(opts cliproxyexecutor.Options, fallback string) string {
|
|
fallback = strings.TrimSpace(fallback)
|
|
if len(opts.Metadata) == 0 {
|
|
return fallback
|
|
}
|
|
raw, ok := opts.Metadata[cliproxyexecutor.RequestedModelMetadataKey]
|
|
if !ok || raw == nil {
|
|
return fallback
|
|
}
|
|
switch value := raw.(type) {
|
|
case string:
|
|
if strings.TrimSpace(value) == "" {
|
|
return fallback
|
|
}
|
|
return strings.TrimSpace(value)
|
|
case []byte:
|
|
if len(value) == 0 {
|
|
return fallback
|
|
}
|
|
return strings.TrimSpace(string(value))
|
|
default:
|
|
return fallback
|
|
}
|
|
}
|
|
|
|
func reasoningEffortFromOptions(opts cliproxyexecutor.Options) string {
|
|
if len(opts.Metadata) == 0 {
|
|
return ""
|
|
}
|
|
raw, ok := opts.Metadata[cliproxyexecutor.ReasoningEffortMetadataKey]
|
|
if !ok || raw == nil {
|
|
return ""
|
|
}
|
|
switch value := raw.(type) {
|
|
case string:
|
|
return strings.TrimSpace(value)
|
|
case []byte:
|
|
return strings.TrimSpace(string(value))
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func serviceTierFromOptions(opts cliproxyexecutor.Options) string {
|
|
if len(opts.Metadata) == 0 {
|
|
return ""
|
|
}
|
|
raw, ok := opts.Metadata[cliproxyexecutor.ServiceTierMetadataKey]
|
|
if !ok || raw == nil {
|
|
return ""
|
|
}
|
|
switch value := raw.(type) {
|
|
case string:
|
|
return strings.TrimSpace(value)
|
|
case []byte:
|
|
return strings.TrimSpace(string(value))
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func pinnedAuthIDFromMetadata(meta map[string]any) string {
|
|
if len(meta) == 0 {
|
|
return ""
|
|
}
|
|
raw, ok := meta[cliproxyexecutor.PinnedAuthMetadataKey]
|
|
if !ok || raw == nil {
|
|
return ""
|
|
}
|
|
switch val := raw.(type) {
|
|
case string:
|
|
return strings.TrimSpace(val)
|
|
case []byte:
|
|
return strings.TrimSpace(string(val))
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func disallowFreeAuthFromMetadata(meta map[string]any) bool {
|
|
if len(meta) == 0 {
|
|
return false
|
|
}
|
|
raw, ok := meta[cliproxyexecutor.DisallowFreeAuthMetadataKey]
|
|
if !ok || raw == nil {
|
|
return false
|
|
}
|
|
switch val := raw.(type) {
|
|
case bool:
|
|
return val
|
|
case string:
|
|
parsed, err := strconv.ParseBool(strings.TrimSpace(val))
|
|
return err == nil && parsed
|
|
case []byte:
|
|
parsed, err := strconv.ParseBool(strings.TrimSpace(string(val)))
|
|
return err == nil && parsed
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func isFreeCodexAuth(auth *Auth) bool {
|
|
if auth == nil || auth.Attributes == nil {
|
|
return false
|
|
}
|
|
if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") {
|
|
return false
|
|
}
|
|
return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free")
|
|
}
|
|
|
|
func publishSelectedAuthMetadata(meta map[string]any, authID string) {
|
|
if len(meta) == 0 {
|
|
return
|
|
}
|
|
authID = strings.TrimSpace(authID)
|
|
if authID == "" {
|
|
return
|
|
}
|
|
meta[cliproxyexecutor.SelectedAuthMetadataKey] = authID
|
|
if callback, ok := meta[cliproxyexecutor.SelectedAuthCallbackMetadataKey].(func(string)); ok && callback != nil {
|
|
callback(authID)
|
|
}
|
|
}
|
|
|
|
func rewriteModelForAuth(model string, auth *Auth) string {
|
|
if auth == nil || model == "" {
|
|
return model
|
|
}
|
|
prefix := strings.TrimSpace(auth.Prefix)
|
|
if prefix == "" {
|
|
return model
|
|
}
|
|
needle := prefix + "/"
|
|
if !strings.HasPrefix(model, needle) {
|
|
return model
|
|
}
|
|
return strings.TrimPrefix(model, needle)
|
|
}
|
|
|
|
func (m *Manager) applyAPIKeyModelAlias(auth *Auth, requestedModel string) string {
|
|
if m == nil || auth == nil {
|
|
return requestedModel
|
|
}
|
|
|
|
if auth.AuthKind() != AuthKindAPIKey {
|
|
return requestedModel
|
|
}
|
|
|
|
requestedModel = strings.TrimSpace(requestedModel)
|
|
if requestedModel == "" {
|
|
return requestedModel
|
|
}
|
|
|
|
// Fast path: lookup per-auth mapping table (keyed by auth.ID).
|
|
if resolved := m.lookupAPIKeyUpstreamModel(auth.ID, requestedModel); resolved != "" {
|
|
return resolved
|
|
}
|
|
|
|
// Slow path: scan config for the matching credential entry and resolve alias.
|
|
// This acts as a safety net if mappings are stale or auth.ID is missing.
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
if cfg == nil {
|
|
cfg = &internalconfig.Config{}
|
|
}
|
|
|
|
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
|
upstreamModel := ""
|
|
switch provider {
|
|
case "gemini":
|
|
upstreamModel = resolveUpstreamModelForGeminiAPIKey(cfg, auth, requestedModel)
|
|
case "gemini-interactions":
|
|
upstreamModel = resolveUpstreamModelForInteractionsAPIKey(cfg, auth, requestedModel)
|
|
case "claude":
|
|
upstreamModel = resolveUpstreamModelForClaudeAPIKey(cfg, auth, requestedModel)
|
|
case "codex":
|
|
upstreamModel = resolveUpstreamModelForCodexAPIKey(cfg, auth, requestedModel)
|
|
case "vertex":
|
|
upstreamModel = resolveUpstreamModelForVertexAPIKey(cfg, auth, requestedModel)
|
|
default:
|
|
upstreamModel = resolveUpstreamModelForOpenAICompatAPIKey(cfg, auth, requestedModel)
|
|
}
|
|
|
|
// Return upstream model if found, otherwise return requested model.
|
|
if upstreamModel != "" {
|
|
return upstreamModel
|
|
}
|
|
return requestedModel
|
|
}
|
|
|
|
// APIKeyConfigEntry is a generic interface for API key configurations.
|
|
type APIKeyConfigEntry interface {
|
|
GetAPIKey() string
|
|
GetBaseURL() string
|
|
}
|
|
|
|
func resolveAPIKeyConfig[T APIKeyConfigEntry](entries []T, auth *Auth) *T {
|
|
if auth == nil || len(entries) == 0 {
|
|
return nil
|
|
}
|
|
attrKey, attrBase := "", ""
|
|
if auth.Attributes != nil {
|
|
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
|
|
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
|
|
}
|
|
for i := range entries {
|
|
entry := &entries[i]
|
|
cfgKey := strings.TrimSpace((*entry).GetAPIKey())
|
|
cfgBase := strings.TrimSpace((*entry).GetBaseURL())
|
|
if attrKey != "" && attrBase != "" {
|
|
if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) {
|
|
return entry
|
|
}
|
|
continue
|
|
}
|
|
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
|
|
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
|
|
return entry
|
|
}
|
|
}
|
|
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
|
|
return entry
|
|
}
|
|
}
|
|
if attrKey != "" {
|
|
for i := range entries {
|
|
entry := &entries[i]
|
|
if strings.EqualFold(strings.TrimSpace((*entry).GetAPIKey()), attrKey) {
|
|
return entry
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func resolveGeminiAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.GeminiKey {
|
|
if cfg == nil {
|
|
return nil
|
|
}
|
|
return resolveAPIKeyConfig(cfg.GeminiKey, auth)
|
|
}
|
|
|
|
func resolveInteractionsAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.GeminiKey {
|
|
if cfg == nil {
|
|
return nil
|
|
}
|
|
return resolveAPIKeyConfig(cfg.InteractionsKey, auth)
|
|
}
|
|
|
|
func resolveClaudeAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.ClaudeKey {
|
|
if cfg == nil {
|
|
return nil
|
|
}
|
|
return resolveAPIKeyConfig(cfg.ClaudeKey, auth)
|
|
}
|
|
|
|
func resolveCodexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.CodexKey {
|
|
if cfg == nil {
|
|
return nil
|
|
}
|
|
return resolveAPIKeyConfig(cfg.CodexKey, auth)
|
|
}
|
|
|
|
func resolveVertexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.VertexCompatKey {
|
|
if cfg == nil {
|
|
return nil
|
|
}
|
|
return resolveAPIKeyConfig(cfg.VertexCompatAPIKey, auth)
|
|
}
|
|
|
|
func resolveUpstreamModelForGeminiAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
|
|
entry := resolveGeminiAPIKeyConfig(cfg, auth)
|
|
if entry == nil {
|
|
return ""
|
|
}
|
|
return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
|
|
}
|
|
|
|
func resolveUpstreamModelForInteractionsAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
|
|
entry := resolveInteractionsAPIKeyConfig(cfg, auth)
|
|
if entry == nil {
|
|
return ""
|
|
}
|
|
return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
|
|
}
|
|
|
|
func resolveUpstreamModelForClaudeAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
|
|
entry := resolveClaudeAPIKeyConfig(cfg, auth)
|
|
if entry == nil {
|
|
return ""
|
|
}
|
|
return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
|
|
}
|
|
|
|
func resolveUpstreamModelForCodexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
|
|
entry := resolveCodexAPIKeyConfig(cfg, auth)
|
|
if entry == nil {
|
|
return ""
|
|
}
|
|
return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
|
|
}
|
|
|
|
func resolveUpstreamModelForVertexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
|
|
entry := resolveVertexAPIKeyConfig(cfg, auth)
|
|
if entry == nil {
|
|
return ""
|
|
}
|
|
return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
|
|
}
|
|
|
|
func resolveUpstreamModelForOpenAICompatAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string {
|
|
providerKey := ""
|
|
compatName := ""
|
|
if auth != nil && len(auth.Attributes) > 0 {
|
|
providerKey = strings.TrimSpace(auth.Attributes["provider_key"])
|
|
compatName = strings.TrimSpace(auth.Attributes["compat_name"])
|
|
}
|
|
if compatName == "" && !strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
|
|
return ""
|
|
}
|
|
entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider)
|
|
if entry == nil {
|
|
return ""
|
|
}
|
|
return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models))
|
|
}
|
|
|
|
type apiKeyModelAliasTable map[string]map[string]string
|
|
|
|
func resolveOpenAICompatConfig(cfg *internalconfig.Config, providerKey, compatName, authProvider string) *internalconfig.OpenAICompatibility {
|
|
if cfg == nil {
|
|
return nil
|
|
}
|
|
candidates := make([]string, 0, 3)
|
|
if v := strings.TrimSpace(compatName); v != "" {
|
|
candidates = append(candidates, v)
|
|
}
|
|
if v := strings.TrimSpace(providerKey); v != "" {
|
|
candidates = append(candidates, v)
|
|
}
|
|
if v := strings.TrimSpace(authProvider); v != "" {
|
|
candidates = append(candidates, v)
|
|
}
|
|
for i := range cfg.OpenAICompatibility {
|
|
compat := &cfg.OpenAICompatibility[i]
|
|
if compat.Disabled {
|
|
continue
|
|
}
|
|
for _, candidate := range candidates {
|
|
if candidate != "" && strings.EqualFold(strings.TrimSpace(candidate), compat.Name) {
|
|
return compat
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func asModelAliasEntries[T interface {
|
|
GetName() string
|
|
GetAlias() string
|
|
GetForceMapping() bool
|
|
}](models []T) []modelAliasEntry {
|
|
if len(models) == 0 {
|
|
return nil
|
|
}
|
|
out := make([]modelAliasEntry, 0, len(models))
|
|
for i := range models {
|
|
out = append(out, models[i])
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (m *Manager) normalizeProviders(providers []string) []string {
|
|
if len(providers) == 0 {
|
|
return nil
|
|
}
|
|
result := make([]string, 0, len(providers))
|
|
seen := make(map[string]struct{}, len(providers))
|
|
for _, provider := range providers {
|
|
p := strings.TrimSpace(strings.ToLower(provider))
|
|
if p == "" {
|
|
continue
|
|
}
|
|
if _, ok := seen[p]; ok {
|
|
continue
|
|
}
|
|
seen[p] = struct{}{}
|
|
result = append(result, p)
|
|
}
|
|
return result
|
|
}
|
|
|
|
// AvailableProviders returns the set of provider keys that currently have at least one
|
|
// registered auth record that is not disabled. It is a best-effort snapshot for routing
|
|
// decisions and does not account for per-model cooldowns or transient runtime availability.
|
|
// Disabled auths (Disabled flag or StatusDisabled) are excluded so routing does not target
|
|
// providers that auth selection would refuse to use, which would otherwise cause execution
|
|
// failures instead of falling back to lower-priority routers.
|
|
func (m *Manager) AvailableProviders() []string {
|
|
if m == nil {
|
|
return nil
|
|
}
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
seen := make(map[string]struct{}, len(m.auths))
|
|
out := make([]string, 0, len(m.auths))
|
|
for _, auth := range m.auths {
|
|
if auth == nil || auth.Disabled || auth.Status == StatusDisabled {
|
|
continue
|
|
}
|
|
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
|
if provider == "" {
|
|
continue
|
|
}
|
|
if _, ok := seen[provider]; ok {
|
|
continue
|
|
}
|
|
seen[provider] = struct{}{}
|
|
out = append(out, provider)
|
|
}
|
|
sort.Strings(out)
|
|
return out
|
|
}
|
|
|
|
// HasProviderAuth reports whether at least one non-disabled auth record is registered for
|
|
// the provider. Disabled auths (Disabled flag or StatusDisabled) are excluded to match the
|
|
// behavior of auth selection, which refuses to pick disabled credentials.
|
|
func (m *Manager) HasProviderAuth(provider string) bool {
|
|
if m == nil {
|
|
return false
|
|
}
|
|
provider = strings.ToLower(strings.TrimSpace(provider))
|
|
if provider == "" {
|
|
return false
|
|
}
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
for _, auth := range m.auths {
|
|
if auth == nil || auth.Disabled || auth.Status == StatusDisabled {
|
|
continue
|
|
}
|
|
if strings.ToLower(strings.TrimSpace(auth.Provider)) == provider {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (m *Manager) retrySettings() (int, int, time.Duration) {
|
|
if m == nil {
|
|
return 0, 0, 0
|
|
}
|
|
return int(m.requestRetry.Load()), int(m.maxRetryCredentials.Load()), time.Duration(m.maxRetryInterval.Load())
|
|
}
|
|
|
|
func (m *Manager) closestCooldownWait(providers []string, model string, attempt int) (time.Duration, bool) {
|
|
if m == nil || len(providers) == 0 {
|
|
return 0, false
|
|
}
|
|
now := time.Now()
|
|
defaultRetry := int(m.requestRetry.Load())
|
|
if defaultRetry < 0 {
|
|
defaultRetry = 0
|
|
}
|
|
providerSet := make(map[string]struct{}, len(providers))
|
|
for i := range providers {
|
|
key := strings.TrimSpace(strings.ToLower(providers[i]))
|
|
if key == "" {
|
|
continue
|
|
}
|
|
providerSet[key] = struct{}{}
|
|
}
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
var (
|
|
found bool
|
|
minWait time.Duration
|
|
)
|
|
for _, auth := range m.auths {
|
|
if auth == nil {
|
|
continue
|
|
}
|
|
providerKey := executorKeyFromAuth(auth)
|
|
if _, ok := providerSet[providerKey]; !ok {
|
|
continue
|
|
}
|
|
effectiveRetry := defaultRetry
|
|
if override, ok := auth.RequestRetryOverride(); ok {
|
|
effectiveRetry = override
|
|
}
|
|
if effectiveRetry < 0 {
|
|
effectiveRetry = 0
|
|
}
|
|
if attempt >= effectiveRetry {
|
|
continue
|
|
}
|
|
checkModel := model
|
|
if strings.TrimSpace(model) != "" {
|
|
checkModel = m.selectionModelForAuth(auth, model)
|
|
}
|
|
blocked, reason, next := isAuthBlockedForModel(auth, checkModel, now)
|
|
if !blocked || next.IsZero() || reason == blockReasonDisabled {
|
|
continue
|
|
}
|
|
wait := next.Sub(now)
|
|
if wait < 0 {
|
|
continue
|
|
}
|
|
if !found || wait < minWait {
|
|
minWait = wait
|
|
found = true
|
|
}
|
|
}
|
|
return minWait, found
|
|
}
|
|
|
|
func (m *Manager) retryAllowed(attempt int, providers []string) bool {
|
|
if m == nil || attempt < 0 || len(providers) == 0 {
|
|
return false
|
|
}
|
|
defaultRetry := int(m.requestRetry.Load())
|
|
if defaultRetry < 0 {
|
|
defaultRetry = 0
|
|
}
|
|
providerSet := make(map[string]struct{}, len(providers))
|
|
for i := range providers {
|
|
key := strings.TrimSpace(strings.ToLower(providers[i]))
|
|
if key == "" {
|
|
continue
|
|
}
|
|
providerSet[key] = struct{}{}
|
|
}
|
|
if len(providerSet) == 0 {
|
|
return false
|
|
}
|
|
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
for _, auth := range m.auths {
|
|
if auth == nil {
|
|
continue
|
|
}
|
|
providerKey := executorKeyFromAuth(auth)
|
|
if _, ok := providerSet[providerKey]; !ok {
|
|
continue
|
|
}
|
|
effectiveRetry := defaultRetry
|
|
if override, ok := auth.RequestRetryOverride(); ok {
|
|
effectiveRetry = override
|
|
}
|
|
if effectiveRetry < 0 {
|
|
effectiveRetry = 0
|
|
}
|
|
if attempt < effectiveRetry {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []string, model string, maxWait time.Duration) (time.Duration, bool) {
|
|
if err == nil {
|
|
return 0, false
|
|
}
|
|
if maxWait <= 0 {
|
|
return 0, false
|
|
}
|
|
status := statusCodeFromError(err)
|
|
if status == http.StatusOK {
|
|
return 0, false
|
|
}
|
|
if isRequestInvalidError(err) {
|
|
return 0, false
|
|
}
|
|
wait, found := m.closestCooldownWait(providers, model, attempt)
|
|
if found {
|
|
if wait > maxWait {
|
|
return 0, false
|
|
}
|
|
return wait, true
|
|
}
|
|
if status != http.StatusTooManyRequests {
|
|
return 0, false
|
|
}
|
|
if !m.retryAllowed(attempt, providers) {
|
|
return 0, false
|
|
}
|
|
retryAfter := retryAfterFromError(err)
|
|
if retryAfter == nil || *retryAfter <= 0 || *retryAfter > maxWait {
|
|
return 0, false
|
|
}
|
|
return *retryAfter, true
|
|
}
|
|
|
|
// cooldownWaitJitterCap bounds the random jitter added to cooldown waits so a
|
|
// long wait is never extended by more than this amount.
|
|
const cooldownWaitJitterCap = 2 * time.Second
|
|
|
|
// jitteredCooldownWait adds a small random delay to a cooldown wait so
|
|
// concurrent requests waiting on the same recovery deadline do not wake in
|
|
// lockstep and stampede the first credential that recovers. The jitter never
|
|
// pushes the total wait past maxWait, which callers have already enforced as
|
|
// the retry ceiling; maxWait <= 0 means no ceiling.
|
|
func jitteredCooldownWait(wait, maxWait time.Duration) time.Duration {
|
|
if wait <= 0 {
|
|
return wait
|
|
}
|
|
jitterRange := wait / 4
|
|
if jitterRange > cooldownWaitJitterCap {
|
|
jitterRange = cooldownWaitJitterCap
|
|
}
|
|
if maxWait > 0 && jitterRange > maxWait-wait {
|
|
jitterRange = maxWait - wait
|
|
}
|
|
if jitterRange <= 0 {
|
|
return wait
|
|
}
|
|
return wait + rand.N(jitterRange)
|
|
}
|
|
|
|
func waitForCooldown(ctx context.Context, wait, maxWait time.Duration) error {
|
|
if wait <= 0 {
|
|
return nil
|
|
}
|
|
timer := time.NewTimer(jitteredCooldownWait(wait, maxWait))
|
|
defer timer.Stop()
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-timer.C:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// MarkResult records an execution result and notifies hooks.
|
|
func (m *Manager) MarkResult(ctx context.Context, result Result) {
|
|
if result.AuthID == "" {
|
|
return
|
|
}
|
|
|
|
shouldResumeModel := false
|
|
shouldSuspendModel := false
|
|
suspendReason := ""
|
|
clearModelQuota := false
|
|
setModelQuota := false
|
|
var authSnapshot *Auth
|
|
cooldownStateChanged := false
|
|
|
|
m.mu.Lock()
|
|
if auth, ok := m.auths[result.AuthID]; ok && auth != nil {
|
|
now := time.Now()
|
|
var cooldownRecordsBefore []CooldownStateRecord
|
|
trackCooldownState := m.cooldownStore != nil
|
|
if trackCooldownState {
|
|
cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now)
|
|
}
|
|
auth.recordRecentRequest(now, result.Success)
|
|
if result.Success {
|
|
auth.Success++
|
|
} else {
|
|
auth.Failed++
|
|
}
|
|
|
|
if result.Success {
|
|
if result.Model != "" {
|
|
state := ensureModelState(auth, result.Model)
|
|
resetModelState(state, now)
|
|
updateAggregatedAvailability(auth, now)
|
|
if !hasModelError(auth, now) {
|
|
auth.LastError = nil
|
|
auth.StatusMessage = ""
|
|
auth.Status = StatusActive
|
|
}
|
|
auth.UpdatedAt = now
|
|
shouldResumeModel = true
|
|
clearModelQuota = true
|
|
} else {
|
|
clearAuthStateOnSuccess(auth, now)
|
|
}
|
|
} else {
|
|
if result.Model != "" {
|
|
if !isRequestScopedNotFoundResultError(result.Error) {
|
|
disableCooling := m.cooldownDisabledForAuth(auth)
|
|
state := ensureModelState(auth, result.Model)
|
|
state.Unavailable = true
|
|
state.Status = StatusError
|
|
state.UpdatedAt = now
|
|
if result.Error != nil {
|
|
state.LastError = cloneError(result.Error)
|
|
state.StatusMessage = result.Error.Message
|
|
auth.LastError = cloneError(result.Error)
|
|
auth.StatusMessage = result.Error.Message
|
|
}
|
|
|
|
statusCode := statusCodeFromResult(result.Error)
|
|
if isModelSupportResultError(result.Error) {
|
|
next := now.Add(12 * time.Hour)
|
|
state.NextRetryAfter = next
|
|
suspendReason = "model_not_supported"
|
|
shouldSuspendModel = true
|
|
} else if isCloudflareChallengeResultError(result.Error) {
|
|
next, backoffLevel := nextCloudflareCooldown(state.Quota.BackoffLevel, disableCooling, now)
|
|
state.NextRetryAfter = next
|
|
state.StatusMessage = "cloudflare challenge"
|
|
if auth.LastError != nil {
|
|
auth.StatusMessage = "cloudflare challenge"
|
|
}
|
|
state.Quota = QuotaState{
|
|
Exceeded: true,
|
|
Reason: "cloudflare challenge",
|
|
NextRecoverAt: next,
|
|
BackoffLevel: backoffLevel,
|
|
}
|
|
} else if isInvalidGrantResultError(result.Error) {
|
|
if disableCooling {
|
|
state.NextRetryAfter = time.Time{}
|
|
} else {
|
|
state.NextRetryAfter = now.Add(30 * time.Minute)
|
|
suspendReason = "invalid_grant"
|
|
shouldSuspendModel = true
|
|
}
|
|
} else {
|
|
switch statusCode {
|
|
case 401:
|
|
if disableCooling {
|
|
state.NextRetryAfter = time.Time{}
|
|
} else {
|
|
next := now.Add(30 * time.Minute)
|
|
state.NextRetryAfter = next
|
|
suspendReason = "unauthorized"
|
|
shouldSuspendModel = true
|
|
}
|
|
case 402, 403:
|
|
if disableCooling {
|
|
state.NextRetryAfter = time.Time{}
|
|
} else {
|
|
next := now.Add(30 * time.Minute)
|
|
state.NextRetryAfter = next
|
|
suspendReason = "payment_required"
|
|
shouldSuspendModel = true
|
|
}
|
|
case 404:
|
|
if disableCooling {
|
|
state.NextRetryAfter = time.Time{}
|
|
} else {
|
|
next := now.Add(12 * time.Hour)
|
|
state.NextRetryAfter = next
|
|
suspendReason = "not_found"
|
|
shouldSuspendModel = true
|
|
}
|
|
case 429:
|
|
var next time.Time
|
|
backoffLevel := state.Quota.BackoffLevel
|
|
if !disableCooling {
|
|
if result.RetryAfter != nil {
|
|
next = now.Add(*result.RetryAfter)
|
|
} else {
|
|
next, backoffLevel = quotaCooldownAfterFailure(state.Quota, now)
|
|
}
|
|
}
|
|
state.NextRetryAfter = next
|
|
state.Quota = QuotaState{
|
|
Exceeded: true,
|
|
Reason: "quota",
|
|
NextRecoverAt: next,
|
|
BackoffLevel: backoffLevel,
|
|
}
|
|
if !disableCooling {
|
|
suspendReason = "quota"
|
|
shouldSuspendModel = true
|
|
setModelQuota = true
|
|
}
|
|
case 408, 500, 502, 503, 504:
|
|
if disableCooling {
|
|
state.NextRetryAfter = time.Time{}
|
|
} else {
|
|
state.NextRetryAfter = nextTransientErrorRetryAfter(now)
|
|
}
|
|
default:
|
|
state.NextRetryAfter = time.Time{}
|
|
}
|
|
}
|
|
|
|
auth.Status = StatusError
|
|
auth.UpdatedAt = now
|
|
updateAggregatedAvailability(auth, now)
|
|
}
|
|
} else {
|
|
disableCooling := m.cooldownDisabledForAuth(auth)
|
|
applyAuthFailureState(auth, result.Error, result.RetryAfter, now, disableCooling)
|
|
}
|
|
}
|
|
|
|
_ = m.persist(ctx, auth)
|
|
authSnapshot = auth.Clone()
|
|
if trackCooldownState {
|
|
cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now)
|
|
cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter)
|
|
}
|
|
}
|
|
m.mu.Unlock()
|
|
if m.scheduler != nil && authSnapshot != nil {
|
|
m.scheduler.upsertAuth(authSnapshot)
|
|
}
|
|
if authSnapshot != nil && cooldownStateChanged {
|
|
m.persistCooldownStates(context.Background())
|
|
}
|
|
|
|
if clearModelQuota && result.Model != "" {
|
|
registry.GetGlobalRegistry().ClearModelQuotaExceeded(result.AuthID, result.Model)
|
|
}
|
|
if setModelQuota && result.Model != "" {
|
|
registry.GetGlobalRegistry().SetModelQuotaExceeded(result.AuthID, result.Model)
|
|
}
|
|
if shouldResumeModel {
|
|
registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, result.Model)
|
|
} else if shouldSuspendModel {
|
|
registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, result.Model, suspendReason)
|
|
}
|
|
|
|
m.hook.OnResult(ctx, result)
|
|
m.publishErrorEvent(result, authSnapshot)
|
|
}
|
|
|
|
func ensureModelState(auth *Auth, model string) *ModelState {
|
|
if auth == nil || model == "" {
|
|
return nil
|
|
}
|
|
if auth.ModelStates == nil {
|
|
auth.ModelStates = make(map[string]*ModelState)
|
|
}
|
|
if state, ok := auth.ModelStates[model]; ok && state != nil {
|
|
return state
|
|
}
|
|
state := &ModelState{Status: StatusActive}
|
|
auth.ModelStates[model] = state
|
|
return state
|
|
}
|
|
|
|
func resetModelState(state *ModelState, now time.Time) {
|
|
if state == nil {
|
|
return
|
|
}
|
|
state.Unavailable = false
|
|
state.Status = StatusActive
|
|
state.StatusMessage = ""
|
|
state.NextRetryAfter = time.Time{}
|
|
state.LastError = nil
|
|
state.Quota = QuotaState{}
|
|
state.UpdatedAt = now
|
|
}
|
|
|
|
func modelStateIsClean(state *ModelState) bool {
|
|
if state == nil {
|
|
return true
|
|
}
|
|
if state.Status != StatusActive {
|
|
return false
|
|
}
|
|
if state.Unavailable || state.StatusMessage != "" || !state.NextRetryAfter.IsZero() || state.LastError != nil {
|
|
return false
|
|
}
|
|
if state.Quota.Exceeded || state.Quota.Reason != "" || !state.Quota.NextRecoverAt.IsZero() || state.Quota.BackoffLevel != 0 {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
func updateAggregatedAvailability(auth *Auth, now time.Time) {
|
|
if auth == nil {
|
|
return
|
|
}
|
|
if len(auth.ModelStates) == 0 {
|
|
clearAggregatedAvailability(auth)
|
|
return
|
|
}
|
|
allUnavailable := true
|
|
earliestRetry := time.Time{}
|
|
quotaExceeded := false
|
|
quotaRecover := time.Time{}
|
|
maxBackoffLevel := 0
|
|
hasState := false
|
|
for _, state := range auth.ModelStates {
|
|
if state == nil {
|
|
continue
|
|
}
|
|
hasState = true
|
|
stateUnavailable := false
|
|
if state.Status == StatusDisabled {
|
|
stateUnavailable = true
|
|
} else if state.Unavailable {
|
|
if state.NextRetryAfter.IsZero() {
|
|
stateUnavailable = false
|
|
} else if state.NextRetryAfter.After(now) {
|
|
stateUnavailable = true
|
|
if earliestRetry.IsZero() || state.NextRetryAfter.Before(earliestRetry) {
|
|
earliestRetry = state.NextRetryAfter
|
|
}
|
|
} else {
|
|
state.Unavailable = false
|
|
state.NextRetryAfter = time.Time{}
|
|
}
|
|
}
|
|
if !stateUnavailable {
|
|
allUnavailable = false
|
|
}
|
|
if state.Quota.Exceeded {
|
|
quotaExceeded = true
|
|
if quotaRecover.IsZero() || (!state.Quota.NextRecoverAt.IsZero() && state.Quota.NextRecoverAt.Before(quotaRecover)) {
|
|
quotaRecover = state.Quota.NextRecoverAt
|
|
}
|
|
if state.Quota.BackoffLevel > maxBackoffLevel {
|
|
maxBackoffLevel = state.Quota.BackoffLevel
|
|
}
|
|
}
|
|
}
|
|
if !hasState {
|
|
clearAggregatedAvailability(auth)
|
|
return
|
|
}
|
|
auth.Unavailable = allUnavailable
|
|
if allUnavailable {
|
|
auth.NextRetryAfter = earliestRetry
|
|
} else {
|
|
auth.NextRetryAfter = time.Time{}
|
|
}
|
|
if quotaExceeded {
|
|
auth.Quota.Exceeded = true
|
|
auth.Quota.Reason = "quota"
|
|
auth.Quota.NextRecoverAt = quotaRecover
|
|
auth.Quota.BackoffLevel = maxBackoffLevel
|
|
} else {
|
|
auth.Quota.Exceeded = false
|
|
auth.Quota.Reason = ""
|
|
auth.Quota.NextRecoverAt = time.Time{}
|
|
auth.Quota.BackoffLevel = 0
|
|
}
|
|
}
|
|
|
|
func clearAggregatedAvailability(auth *Auth) {
|
|
if auth == nil {
|
|
return
|
|
}
|
|
auth.Unavailable = false
|
|
auth.NextRetryAfter = time.Time{}
|
|
auth.Quota = QuotaState{}
|
|
}
|
|
|
|
func hasModelError(auth *Auth, now time.Time) bool {
|
|
if auth == nil || len(auth.ModelStates) == 0 {
|
|
return false
|
|
}
|
|
for _, state := range auth.ModelStates {
|
|
if state == nil {
|
|
continue
|
|
}
|
|
if state.LastError != nil {
|
|
return true
|
|
}
|
|
if state.Status == StatusError {
|
|
if state.Unavailable && (state.NextRetryAfter.IsZero() || state.NextRetryAfter.After(now)) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func clearAuthStateOnSuccess(auth *Auth, now time.Time) {
|
|
if auth == nil {
|
|
return
|
|
}
|
|
auth.Unavailable = false
|
|
auth.Status = StatusActive
|
|
auth.StatusMessage = ""
|
|
auth.Quota.Exceeded = false
|
|
auth.Quota.Reason = ""
|
|
auth.Quota.NextRecoverAt = time.Time{}
|
|
auth.Quota.BackoffLevel = 0
|
|
auth.LastError = nil
|
|
auth.NextRetryAfter = time.Time{}
|
|
auth.UpdatedAt = now
|
|
}
|
|
|
|
func cloneError(err *Error) *Error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
return &Error{
|
|
Code: err.Code,
|
|
Message: err.Message,
|
|
Retryable: err.Retryable,
|
|
HTTPStatus: err.HTTPStatus,
|
|
}
|
|
}
|
|
|
|
func errorString(err error) string {
|
|
if err == nil {
|
|
return ""
|
|
}
|
|
return err.Error()
|
|
}
|
|
|
|
func statusCodeFromError(err error) int {
|
|
if err == nil {
|
|
return 0
|
|
}
|
|
type statusCoder interface {
|
|
StatusCode() int
|
|
}
|
|
var sc statusCoder
|
|
if errors.As(err, &sc) && sc != nil {
|
|
return sc.StatusCode()
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func isUnauthorizedError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
if statusCodeFromError(err) == http.StatusUnauthorized {
|
|
return true
|
|
}
|
|
raw := strings.ToLower(err.Error())
|
|
return strings.Contains(raw, "status 401") || strings.Contains(raw, "401 unauthorized")
|
|
}
|
|
|
|
func hasUnauthorizedAuthFailure(auth *Auth) bool {
|
|
if auth == nil || auth.LastError == nil {
|
|
return false
|
|
}
|
|
return auth.LastError.StatusCode() == http.StatusUnauthorized || strings.EqualFold(auth.LastError.Code, "unauthorized")
|
|
}
|
|
|
|
func refreshErrorFromError(err error) *Error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
statusCode := statusCodeFromError(err)
|
|
if statusCode == 0 && isUnauthorizedError(err) {
|
|
statusCode = http.StatusUnauthorized
|
|
}
|
|
authErr := &Error{Message: err.Error(), HTTPStatus: statusCode}
|
|
if statusCode == http.StatusUnauthorized {
|
|
authErr.Code = "unauthorized"
|
|
authErr.Retryable = false
|
|
}
|
|
return authErr
|
|
}
|
|
|
|
func retryAfterFromError(err error) *time.Duration {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
type retryAfterProvider interface {
|
|
RetryAfter() *time.Duration
|
|
}
|
|
rap, ok := err.(retryAfterProvider)
|
|
if !ok || rap == nil {
|
|
return nil
|
|
}
|
|
retryAfter := rap.RetryAfter()
|
|
if retryAfter == nil {
|
|
return nil
|
|
}
|
|
value := *retryAfter
|
|
return &value
|
|
}
|
|
|
|
func statusCodeFromResult(err *Error) int {
|
|
if err == nil {
|
|
return 0
|
|
}
|
|
return err.StatusCode()
|
|
}
|
|
|
|
func isModelSupportErrorMessage(message string) bool {
|
|
lower := strings.ToLower(strings.TrimSpace(message))
|
|
if lower == "" {
|
|
return false
|
|
}
|
|
patterns := [...]string{
|
|
"model_not_supported",
|
|
"requested model is not supported",
|
|
"requested model is unsupported",
|
|
"requested model is unavailable",
|
|
"model is not supported",
|
|
"model not supported",
|
|
"unsupported model",
|
|
"model unavailable",
|
|
"not available for your plan",
|
|
"not available for your account",
|
|
}
|
|
for _, pattern := range patterns {
|
|
if strings.Contains(lower, pattern) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func isModelSupportError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
status := statusCodeFromError(err)
|
|
if status != http.StatusBadRequest && status != http.StatusUnprocessableEntity {
|
|
return false
|
|
}
|
|
return isModelSupportErrorMessage(err.Error())
|
|
}
|
|
|
|
func isInvalidGrantErrorMessage(message string) bool {
|
|
return strings.Contains(strings.ToLower(message), "invalid_grant")
|
|
}
|
|
|
|
func isInvalidGrantError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
status := statusCodeFromError(err)
|
|
if status != http.StatusBadRequest && status != http.StatusUnauthorized {
|
|
return false
|
|
}
|
|
return isInvalidGrantErrorMessage(err.Error())
|
|
}
|
|
|
|
func isInvalidGrantResultError(err *Error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
status := statusCodeFromResult(err)
|
|
if status != http.StatusBadRequest && status != http.StatusUnauthorized {
|
|
return false
|
|
}
|
|
return isInvalidGrantErrorMessage(err.Code) || isInvalidGrantErrorMessage(err.Message)
|
|
}
|
|
|
|
func isModelSupportResultError(err *Error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
status := statusCodeFromResult(err)
|
|
if status != http.StatusBadRequest && status != http.StatusUnprocessableEntity {
|
|
return false
|
|
}
|
|
return isModelSupportErrorMessage(err.Message)
|
|
}
|
|
|
|
func isCloudflareChallengeErrorMessage(message string) bool {
|
|
lower := strings.ToLower(strings.TrimSpace(message))
|
|
return strings.Contains(lower, "challenge-platform") ||
|
|
strings.Contains(lower, "cf-mitigated") ||
|
|
strings.Contains(lower, "cloudflare challenge") ||
|
|
(strings.Contains(lower, "cloudflare") && strings.Contains(lower, "<html"))
|
|
}
|
|
|
|
func isCloudflareChallengeError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
return isCloudflareChallengeErrorMessage(err.Error())
|
|
}
|
|
|
|
func isCloudflareChallengeResultError(err *Error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
return isCloudflareChallengeErrorMessage(err.Message)
|
|
}
|
|
|
|
func nextCloudflareCooldown(backoffLevel int, disableCooling bool, now time.Time) (time.Time, int) {
|
|
var next time.Time
|
|
if !disableCooling {
|
|
cooldown, nextLevel := nextQuotaCooldown(backoffLevel, disableCooling)
|
|
if cooldown < 10*time.Second {
|
|
cooldown = 10 * time.Second
|
|
}
|
|
if cooldown > 0 {
|
|
next = now.Add(cooldown)
|
|
}
|
|
backoffLevel = nextLevel
|
|
}
|
|
return next, backoffLevel
|
|
}
|
|
func isRequestScopedNotFoundMessage(message string) bool {
|
|
if message == "" {
|
|
return false
|
|
}
|
|
lower := strings.ToLower(message)
|
|
return strings.Contains(lower, "item with id") &&
|
|
strings.Contains(lower, "not found") &&
|
|
strings.Contains(lower, "items are not persisted when `store` is set to false")
|
|
}
|
|
|
|
func isRequestScopedNotFoundResultError(err *Error) bool {
|
|
if err == nil || statusCodeFromResult(err) != http.StatusNotFound {
|
|
return false
|
|
}
|
|
return isRequestScopedNotFoundMessage(err.Message)
|
|
}
|
|
|
|
// isRequestInvalidError returns true if the error represents a client request
|
|
// error that should not be retried. Specifically, it treats 400 responses with
|
|
// "invalid_request_error", request-scoped 404 item misses caused by `store=false`,
|
|
// and all 422 responses as request-shape failures, where switching auths or
|
|
// pooled upstream models will not help. Model-support errors are excluded so
|
|
// routing can fall through to another auth or upstream.
|
|
func isRequestInvalidError(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
if isCloudflareChallengeError(err) {
|
|
return false
|
|
}
|
|
if isInvalidGrantError(err) {
|
|
return false
|
|
}
|
|
if isModelSupportError(err) {
|
|
return false
|
|
}
|
|
status := statusCodeFromError(err)
|
|
switch status {
|
|
case http.StatusBadRequest:
|
|
msg := err.Error()
|
|
return strings.Contains(msg, "invalid_request_error") ||
|
|
strings.Contains(msg, "INVALID_ARGUMENT") ||
|
|
strings.Contains(msg, "FAILED_PRECONDITION")
|
|
case http.StatusNotFound:
|
|
return isRequestScopedNotFoundMessage(err.Error())
|
|
case http.StatusUnprocessableEntity:
|
|
return true
|
|
case http.StatusInternalServerError:
|
|
msg := err.Error()
|
|
return strings.Contains(msg, "\"status\":\"UNKNOWN\"") ||
|
|
strings.Contains(msg, "\"status\": \"UNKNOWN\"")
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Duration, now time.Time, disableCooling bool) {
|
|
if auth == nil {
|
|
return
|
|
}
|
|
if isRequestScopedNotFoundResultError(resultErr) {
|
|
return
|
|
}
|
|
auth.Unavailable = true
|
|
auth.Status = StatusError
|
|
auth.UpdatedAt = now
|
|
if resultErr != nil {
|
|
auth.LastError = cloneError(resultErr)
|
|
if resultErr.Message != "" {
|
|
auth.StatusMessage = resultErr.Message
|
|
}
|
|
}
|
|
statusCode := statusCodeFromResult(resultErr)
|
|
if isCloudflareChallengeResultError(resultErr) {
|
|
auth.StatusMessage = "cloudflare challenge"
|
|
next, backoffLevel := nextCloudflareCooldown(auth.Quota.BackoffLevel, disableCooling, now)
|
|
auth.Quota = QuotaState{
|
|
Exceeded: true,
|
|
Reason: "cloudflare challenge",
|
|
NextRecoverAt: next,
|
|
BackoffLevel: backoffLevel,
|
|
}
|
|
auth.NextRetryAfter = next
|
|
return
|
|
}
|
|
if isInvalidGrantResultError(resultErr) {
|
|
auth.StatusMessage = "invalid_grant"
|
|
if disableCooling {
|
|
auth.NextRetryAfter = time.Time{}
|
|
} else {
|
|
auth.NextRetryAfter = now.Add(30 * time.Minute)
|
|
}
|
|
return
|
|
}
|
|
switch statusCode {
|
|
case 401:
|
|
auth.StatusMessage = "unauthorized"
|
|
if disableCooling {
|
|
auth.NextRetryAfter = time.Time{}
|
|
} else {
|
|
auth.NextRetryAfter = now.Add(30 * time.Minute)
|
|
}
|
|
case 402, 403:
|
|
auth.StatusMessage = "payment_required"
|
|
if disableCooling {
|
|
auth.NextRetryAfter = time.Time{}
|
|
} else {
|
|
auth.NextRetryAfter = now.Add(30 * time.Minute)
|
|
}
|
|
case 404:
|
|
auth.StatusMessage = "not_found"
|
|
if disableCooling {
|
|
auth.NextRetryAfter = time.Time{}
|
|
} else {
|
|
auth.NextRetryAfter = now.Add(12 * time.Hour)
|
|
}
|
|
case 429:
|
|
auth.StatusMessage = "quota exhausted"
|
|
auth.Quota.Exceeded = true
|
|
auth.Quota.Reason = "quota"
|
|
var next time.Time
|
|
if !disableCooling {
|
|
if retryAfter != nil {
|
|
next = now.Add(*retryAfter)
|
|
} else {
|
|
next, auth.Quota.BackoffLevel = quotaCooldownAfterFailure(auth.Quota, now)
|
|
}
|
|
}
|
|
auth.Quota.NextRecoverAt = next
|
|
auth.NextRetryAfter = next
|
|
case 408, 500, 502, 503, 504:
|
|
auth.StatusMessage = "transient upstream error"
|
|
if disableCooling {
|
|
auth.NextRetryAfter = time.Time{}
|
|
} else {
|
|
auth.NextRetryAfter = nextTransientErrorRetryAfter(now)
|
|
}
|
|
default:
|
|
if auth.StatusMessage == "" {
|
|
auth.StatusMessage = "request failed"
|
|
}
|
|
}
|
|
}
|
|
|
|
// quotaCooldownAfterFailure returns the recovery deadline and backoff level for
|
|
// a quota failure observed at now. Failures that land while a previous quota
|
|
// window is still open reuse that window instead of escalating, so a burst of
|
|
// concurrent in-flight failures advances the backoff ladder at most once per
|
|
// window.
|
|
func quotaCooldownAfterFailure(quota QuotaState, now time.Time) (time.Time, int) {
|
|
if quota.NextRecoverAt.After(now) {
|
|
return quota.NextRecoverAt, quota.BackoffLevel
|
|
}
|
|
cooldown, nextLevel := nextQuotaCooldown(quota.BackoffLevel, false)
|
|
var next time.Time
|
|
if cooldown > 0 {
|
|
next = now.Add(cooldown)
|
|
}
|
|
return next, nextLevel
|
|
}
|
|
|
|
// nextQuotaCooldown returns the next cooldown duration and updated backoff level for repeated quota errors.
|
|
func nextQuotaCooldown(prevLevel int, disableCooling bool) (time.Duration, int) {
|
|
if prevLevel < 0 {
|
|
prevLevel = 0
|
|
}
|
|
if disableCooling {
|
|
return 0, prevLevel
|
|
}
|
|
cooldown := quotaBackoffBase * time.Duration(1<<prevLevel)
|
|
if cooldown < quotaBackoffBase {
|
|
cooldown = quotaBackoffBase
|
|
}
|
|
if cooldown >= quotaBackoffMax {
|
|
return quotaBackoffMax, prevLevel
|
|
}
|
|
return cooldown, prevLevel + 1
|
|
}
|
|
|
|
// List returns all auth entries currently known by the manager.
|
|
func (m *Manager) List() []*Auth {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
list := make([]*Auth, 0, len(m.auths))
|
|
for _, auth := range m.auths {
|
|
list = append(list, auth.Clone())
|
|
}
|
|
return list
|
|
}
|
|
|
|
// GetByID retrieves an auth entry by its ID.
|
|
|
|
func (m *Manager) GetByID(id string) (*Auth, bool) {
|
|
if id == "" {
|
|
return nil, false
|
|
}
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
auth, ok := m.auths[id]
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
return auth.Clone(), true
|
|
}
|
|
|
|
// GetExecutionSessionAuthByID retrieves a Home runtime auth scoped to an execution session.
|
|
func (m *Manager) GetExecutionSessionAuthByID(sessionID string, authID string) (*Auth, bool) {
|
|
sessionID = strings.TrimSpace(sessionID)
|
|
authID = strings.TrimSpace(authID)
|
|
if m == nil || sessionID == "" || authID == "" {
|
|
return nil, false
|
|
}
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
sessionAuths := m.homeRuntimeAuths[sessionID]
|
|
auth := sessionAuths[authID]
|
|
if auth == nil {
|
|
return nil, false
|
|
}
|
|
return auth.Clone(), true
|
|
}
|
|
|
|
// Executor returns the registered provider executor for a provider key.
|
|
func (m *Manager) Executor(provider string) (ProviderExecutor, bool) {
|
|
if m == nil {
|
|
return nil, false
|
|
}
|
|
provider = strings.TrimSpace(provider)
|
|
if provider == "" {
|
|
return nil, false
|
|
}
|
|
|
|
m.mu.RLock()
|
|
executor, okExecutor := m.executors[provider]
|
|
if !okExecutor {
|
|
lowerProvider := strings.ToLower(provider)
|
|
if lowerProvider != provider {
|
|
executor, okExecutor = m.executors[lowerProvider]
|
|
}
|
|
}
|
|
m.mu.RUnlock()
|
|
|
|
if !okExecutor || executor == nil {
|
|
return nil, false
|
|
}
|
|
return executor, true
|
|
}
|
|
|
|
// CloseExecutionSession asks all registered executors to release the supplied execution session.
|
|
func (m *Manager) CloseExecutionSession(sessionID string) {
|
|
sessionID = strings.TrimSpace(sessionID)
|
|
if m == nil || sessionID == "" {
|
|
return
|
|
}
|
|
|
|
m.mu.Lock()
|
|
if sessionID == CloseAllExecutionSessionsID {
|
|
m.clearHomeRuntimeAuthsLocked()
|
|
} else {
|
|
m.clearHomeRuntimeAuthsForSessionLocked(sessionID)
|
|
}
|
|
executors := make([]ProviderExecutor, 0, len(m.executors))
|
|
for _, exec := range m.executors {
|
|
executors = append(executors, exec)
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
for i := range executors {
|
|
if closer, ok := executors[i].(ExecutionSessionCloser); ok && closer != nil {
|
|
closer.CloseExecutionSession(sessionID)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *Manager) useSchedulerFastPath() bool {
|
|
if m == nil || m.scheduler == nil {
|
|
return false
|
|
}
|
|
return isBuiltInSelector(m.selector)
|
|
}
|
|
|
|
func shouldRetrySchedulerPick(err error) bool {
|
|
if err == nil {
|
|
return false
|
|
}
|
|
var cooldownErr *modelCooldownError
|
|
if errors.As(err, &cooldownErr) {
|
|
return true
|
|
}
|
|
var authErr *Error
|
|
if !errors.As(err, &authErr) || authErr == nil {
|
|
return false
|
|
}
|
|
return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable"
|
|
}
|
|
|
|
func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) bool {
|
|
if auth == nil || strings.TrimSpace(routeModel) == "" {
|
|
return false
|
|
}
|
|
return m.selectionModelKeyForAuth(auth, routeModel) != canonicalModelKey(routeModel)
|
|
}
|
|
|
|
func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) {
|
|
if m.HomeEnabled() {
|
|
auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried)
|
|
return auth, exec, err
|
|
}
|
|
|
|
pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata)
|
|
disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata)
|
|
|
|
m.mu.RLock()
|
|
selector := m.selector
|
|
pluginScheduler := m.pluginScheduler
|
|
executor, okExecutor := m.executors[provider]
|
|
if !okExecutor {
|
|
m.mu.RUnlock()
|
|
return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"}
|
|
}
|
|
candidates := make([]*Auth, 0, len(m.auths))
|
|
modelKey := strings.TrimSpace(model)
|
|
// Always use base model name (without thinking suffix) for auth matching.
|
|
if modelKey != "" {
|
|
parsed := thinking.ParseSuffix(modelKey)
|
|
if parsed.ModelName != "" {
|
|
modelKey = strings.TrimSpace(parsed.ModelName)
|
|
}
|
|
}
|
|
registryRef := registry.GetGlobalRegistry()
|
|
for _, candidate := range m.auths {
|
|
if candidate == nil || executorKeyFromAuth(candidate) != provider || candidate.Disabled {
|
|
continue
|
|
}
|
|
if pinnedAuthID != "" && candidate.ID != pinnedAuthID {
|
|
continue
|
|
}
|
|
if disallowFreeAuth && isFreeCodexAuth(candidate) {
|
|
continue
|
|
}
|
|
if _, used := tried[candidate.ID]; used {
|
|
continue
|
|
}
|
|
if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) {
|
|
continue
|
|
}
|
|
candidates = append(candidates, candidate)
|
|
}
|
|
if len(candidates) == 0 {
|
|
m.mu.RUnlock()
|
|
return nil, nil, &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
available, errAvailable := m.availableAuthsForRouteModel(candidates, provider, model, time.Now())
|
|
if errAvailable != nil {
|
|
m.mu.RUnlock()
|
|
return nil, nil, errAvailable
|
|
}
|
|
available = cloneAuthSlice(available)
|
|
m.mu.RUnlock()
|
|
|
|
selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, provider, []string{provider}, model, opts, tried, available)
|
|
if errPick != nil {
|
|
return nil, nil, errPick
|
|
}
|
|
if !handled {
|
|
selected, errPick = selector.Pick(ctx, provider, selectionArgForSelector(selector, model), opts, available)
|
|
if errPick != nil {
|
|
return nil, nil, errPick
|
|
}
|
|
}
|
|
if selected == nil {
|
|
return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"}
|
|
}
|
|
authCopy := selected.Clone()
|
|
if !selected.indexAssigned {
|
|
m.mu.Lock()
|
|
if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
|
|
current.EnsureIndex()
|
|
authCopy = current.Clone()
|
|
}
|
|
m.mu.Unlock()
|
|
}
|
|
return authCopy, executor, nil
|
|
}
|
|
|
|
// SelectAuth selects one credential through the configured scheduling strategy.
|
|
// It does not execute or alter the selected credential's result state.
|
|
func (m *Manager) SelectAuth(ctx context.Context, provider, model string, opts cliproxyexecutor.Options) (*Auth, error) {
|
|
selected, _, errPick := m.pickNext(ctx, provider, model, opts, nil)
|
|
if errPick != nil {
|
|
return nil, errPick
|
|
}
|
|
return selected, nil
|
|
}
|
|
|
|
func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) {
|
|
if m.HomeEnabled() {
|
|
auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried)
|
|
return auth, exec, err
|
|
}
|
|
|
|
if m.hasPluginScheduler() || !m.useSchedulerFastPath() {
|
|
return m.pickNextLegacy(ctx, provider, model, opts, tried)
|
|
}
|
|
if strings.TrimSpace(model) != "" {
|
|
m.mu.RLock()
|
|
for _, candidate := range m.auths {
|
|
if candidate == nil || executorKeyFromAuth(candidate) != provider || candidate.Disabled {
|
|
continue
|
|
}
|
|
if _, used := tried[candidate.ID]; used {
|
|
continue
|
|
}
|
|
if m.routeAwareSelectionRequired(candidate, model) {
|
|
m.mu.RUnlock()
|
|
return m.pickNextLegacy(ctx, provider, model, opts, tried)
|
|
}
|
|
}
|
|
m.mu.RUnlock()
|
|
}
|
|
executor, okExecutor := m.Executor(provider)
|
|
if !okExecutor {
|
|
return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"}
|
|
}
|
|
disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata)
|
|
for {
|
|
selected, errPick := m.scheduler.pickSingle(ctx, provider, model, opts, tried)
|
|
if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) {
|
|
m.syncScheduler()
|
|
selected, errPick = m.scheduler.pickSingle(ctx, provider, model, opts, tried)
|
|
}
|
|
if errPick != nil {
|
|
return nil, nil, errPick
|
|
}
|
|
if selected == nil {
|
|
return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"}
|
|
}
|
|
if disallowFreeAuth && isFreeCodexAuth(selected) {
|
|
if tried == nil {
|
|
tried = make(map[string]struct{})
|
|
}
|
|
tried[selected.ID] = struct{}{}
|
|
continue
|
|
}
|
|
authCopy := selected.Clone()
|
|
if !selected.indexAssigned {
|
|
m.mu.Lock()
|
|
if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
|
|
current.EnsureIndex()
|
|
authCopy = current.Clone()
|
|
}
|
|
m.mu.Unlock()
|
|
}
|
|
return authCopy, executor, nil
|
|
}
|
|
}
|
|
|
|
func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) {
|
|
if m.HomeEnabled() {
|
|
return m.pickNextViaHome(ctx, model, opts, tried)
|
|
}
|
|
|
|
pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata)
|
|
disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata)
|
|
|
|
providerSet := make(map[string]struct{}, len(providers))
|
|
for _, provider := range providers {
|
|
p := strings.TrimSpace(strings.ToLower(provider))
|
|
if p == "" {
|
|
continue
|
|
}
|
|
providerSet[p] = struct{}{}
|
|
}
|
|
if len(providerSet) == 0 {
|
|
return nil, nil, "", &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
|
}
|
|
|
|
m.mu.RLock()
|
|
selector := m.selector
|
|
pluginScheduler := m.pluginScheduler
|
|
candidates := make([]*Auth, 0, len(m.auths))
|
|
modelKey := strings.TrimSpace(model)
|
|
// Always use base model name (without thinking suffix) for auth matching.
|
|
if modelKey != "" {
|
|
parsed := thinking.ParseSuffix(modelKey)
|
|
if parsed.ModelName != "" {
|
|
modelKey = strings.TrimSpace(parsed.ModelName)
|
|
}
|
|
}
|
|
registryRef := registry.GetGlobalRegistry()
|
|
for _, candidate := range m.auths {
|
|
if candidate == nil || candidate.Disabled {
|
|
continue
|
|
}
|
|
if pinnedAuthID != "" && candidate.ID != pinnedAuthID {
|
|
continue
|
|
}
|
|
if disallowFreeAuth && isFreeCodexAuth(candidate) {
|
|
continue
|
|
}
|
|
providerKey := executorKeyFromAuth(candidate)
|
|
if providerKey == "" {
|
|
continue
|
|
}
|
|
if _, ok := providerSet[providerKey]; !ok {
|
|
continue
|
|
}
|
|
if _, used := tried[candidate.ID]; used {
|
|
continue
|
|
}
|
|
if _, ok := m.executors[providerKey]; !ok {
|
|
continue
|
|
}
|
|
if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) {
|
|
continue
|
|
}
|
|
candidates = append(candidates, candidate)
|
|
}
|
|
if len(candidates) == 0 {
|
|
m.mu.RUnlock()
|
|
return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
available, errAvailable := m.availableAuthsForRouteModel(candidates, "mixed", model, time.Now())
|
|
if errAvailable != nil {
|
|
m.mu.RUnlock()
|
|
return nil, nil, "", errAvailable
|
|
}
|
|
available = cloneAuthSlice(available)
|
|
m.mu.RUnlock()
|
|
|
|
selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, "mixed", providers, model, opts, tried, available)
|
|
if errPick != nil {
|
|
return nil, nil, "", errPick
|
|
}
|
|
if !handled {
|
|
selected, errPick = selector.Pick(ctx, "mixed", selectionArgForSelector(selector, model), opts, available)
|
|
if errPick != nil {
|
|
return nil, nil, "", errPick
|
|
}
|
|
}
|
|
if selected == nil {
|
|
return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"}
|
|
}
|
|
providerKey := executorKeyFromAuth(selected)
|
|
executor, okExecutor := m.Executor(providerKey)
|
|
if !okExecutor {
|
|
return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"}
|
|
}
|
|
authCopy := selected.Clone()
|
|
if !selected.indexAssigned {
|
|
m.mu.Lock()
|
|
if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
|
|
current.EnsureIndex()
|
|
authCopy = current.Clone()
|
|
}
|
|
m.mu.Unlock()
|
|
}
|
|
return authCopy, executor, providerKey, nil
|
|
}
|
|
|
|
func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) {
|
|
if m.HomeEnabled() {
|
|
return m.pickNextViaHome(ctx, model, opts, tried)
|
|
}
|
|
|
|
if m.hasPluginScheduler() || !m.useSchedulerFastPath() {
|
|
return m.pickNextMixedLegacy(ctx, providers, model, opts, tried)
|
|
}
|
|
|
|
eligibleProviders := make([]string, 0, len(providers))
|
|
seenProviders := make(map[string]struct{}, len(providers))
|
|
for _, provider := range providers {
|
|
providerKey := strings.TrimSpace(strings.ToLower(provider))
|
|
if providerKey == "" {
|
|
continue
|
|
}
|
|
if _, seen := seenProviders[providerKey]; seen {
|
|
continue
|
|
}
|
|
if _, okExecutor := m.Executor(providerKey); !okExecutor {
|
|
continue
|
|
}
|
|
seenProviders[providerKey] = struct{}{}
|
|
eligibleProviders = append(eligibleProviders, providerKey)
|
|
}
|
|
if len(eligibleProviders) == 0 {
|
|
return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
if strings.TrimSpace(model) != "" {
|
|
providerSet := make(map[string]struct{}, len(eligibleProviders))
|
|
for _, providerKey := range eligibleProviders {
|
|
providerSet[providerKey] = struct{}{}
|
|
}
|
|
m.mu.RLock()
|
|
for _, candidate := range m.auths {
|
|
if candidate == nil || candidate.Disabled {
|
|
continue
|
|
}
|
|
if _, ok := providerSet[executorKeyFromAuth(candidate)]; !ok {
|
|
continue
|
|
}
|
|
if _, used := tried[candidate.ID]; used {
|
|
continue
|
|
}
|
|
if m.routeAwareSelectionRequired(candidate, model) {
|
|
m.mu.RUnlock()
|
|
return m.pickNextMixedLegacy(ctx, providers, model, opts, tried)
|
|
}
|
|
}
|
|
m.mu.RUnlock()
|
|
}
|
|
|
|
disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata)
|
|
for {
|
|
selected, providerKey, errPick := m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried)
|
|
if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) {
|
|
m.syncScheduler()
|
|
selected, providerKey, errPick = m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried)
|
|
}
|
|
if errPick != nil {
|
|
return nil, nil, "", errPick
|
|
}
|
|
if selected == nil {
|
|
return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"}
|
|
}
|
|
if disallowFreeAuth && isFreeCodexAuth(selected) {
|
|
if tried == nil {
|
|
tried = make(map[string]struct{})
|
|
}
|
|
tried[selected.ID] = struct{}{}
|
|
continue
|
|
}
|
|
executor, okExecutor := m.Executor(providerKey)
|
|
if !okExecutor {
|
|
return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"}
|
|
}
|
|
authCopy := selected.Clone()
|
|
if !selected.indexAssigned {
|
|
m.mu.Lock()
|
|
if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
|
|
current.EnsureIndex()
|
|
authCopy = current.Clone()
|
|
}
|
|
m.mu.Unlock()
|
|
}
|
|
return authCopy, executor, providerKey, nil
|
|
}
|
|
}
|
|
|
|
type homeErrorEnvelope struct {
|
|
Error *homeErrorDetail `json:"error"`
|
|
}
|
|
|
|
type homeErrorDetail struct {
|
|
Type string `json:"type"`
|
|
Message string `json:"message"`
|
|
Code string `json:"code,omitempty"`
|
|
}
|
|
|
|
const (
|
|
homeUpstreamModelAttributeKey = "home_upstream_model"
|
|
homeRequestRetryExceededErrorCode = "request_retry_exceeded"
|
|
)
|
|
|
|
func isHomeRequestRetryExceededError(err error) bool {
|
|
var authErr *Error
|
|
if !errors.As(err, &authErr) || authErr == nil {
|
|
return false
|
|
}
|
|
return strings.EqualFold(strings.TrimSpace(authErr.Code), homeRequestRetryExceededErrorCode)
|
|
}
|
|
|
|
func shouldReturnLastErrorOnPickFailure(homeMode bool, lastErr error, errPick error) bool {
|
|
if lastErr == nil {
|
|
return false
|
|
}
|
|
if !homeMode {
|
|
return true
|
|
}
|
|
return isHomeRequestRetryExceededError(errPick)
|
|
}
|
|
|
|
func homeAuthAlreadyTried(tried map[string]struct{}, authID string) bool {
|
|
authID = strings.TrimSpace(authID)
|
|
if authID == "" || len(tried) == 0 {
|
|
return false
|
|
}
|
|
_, ok := tried[authID]
|
|
return ok
|
|
}
|
|
|
|
func repeatedHomeAuthError() *Error {
|
|
return &Error{
|
|
Code: homeRequestRetryExceededErrorCode,
|
|
Message: "home returned a previously tried auth",
|
|
HTTPStatus: http.StatusServiceUnavailable,
|
|
}
|
|
}
|
|
|
|
type homeAuthDispatchResponse struct {
|
|
Model string `json:"model"`
|
|
Provider string `json:"provider"`
|
|
AuthIndex string `json:"auth_index"`
|
|
UserAPIKey string `json:"user_api_key"`
|
|
Auth Auth `json:"auth"`
|
|
}
|
|
|
|
type homeAuthDispatcher interface {
|
|
HeartbeatOK() bool
|
|
RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error)
|
|
}
|
|
|
|
var currentHomeDispatcher = func() homeAuthDispatcher {
|
|
return home.Current()
|
|
}
|
|
|
|
func setHomeUserAPIKeyOnGinContext(ctx context.Context, apiKey string) {
|
|
apiKey = strings.TrimSpace(apiKey)
|
|
if apiKey == "" || ctx == nil {
|
|
return
|
|
}
|
|
ginCtx, ok := ctx.Value("gin").(interface{ Set(string, any) })
|
|
if !ok || ginCtx == nil {
|
|
return
|
|
}
|
|
ginCtx.Set("userApiKey", apiKey)
|
|
}
|
|
|
|
func homeDispatchHeaders(ctx context.Context, headers http.Header) http.Header {
|
|
apiKey, ok := homeQueryCredentialFromContext(ctx)
|
|
if !ok {
|
|
return headers
|
|
}
|
|
out := headers.Clone()
|
|
if out == nil {
|
|
out = http.Header{}
|
|
}
|
|
if out.Get("Authorization") != "" || out.Get("X-Goog-Api-Key") != "" || out.Get("X-Api-Key") != "" {
|
|
return out
|
|
}
|
|
out.Set("X-Goog-Api-Key", apiKey)
|
|
return out
|
|
}
|
|
|
|
func homeQueryCredentialFromContext(ctx context.Context) (string, bool) {
|
|
if ctx == nil {
|
|
return "", false
|
|
}
|
|
if queryCtx, ok := ctx.Value("gin").(interface{ Query(string) string }); ok && queryCtx != nil {
|
|
if apiKey := strings.TrimSpace(queryCtx.Query("key")); apiKey != "" {
|
|
return apiKey, true
|
|
}
|
|
if apiKey := strings.TrimSpace(queryCtx.Query("auth_token")); apiKey != "" {
|
|
return apiKey, true
|
|
}
|
|
}
|
|
ginCtx, ok := ctx.Value("gin").(interface{ Get(string) (any, bool) })
|
|
if !ok || ginCtx == nil {
|
|
return "", false
|
|
}
|
|
rawMetadata, ok := ginCtx.Get("accessMetadata")
|
|
if !ok {
|
|
return "", false
|
|
}
|
|
source := accessMetadataSource(rawMetadata)
|
|
if source != "query-key" && source != "query-auth-token" {
|
|
return "", false
|
|
}
|
|
rawAPIKey, ok := ginCtx.Get("userApiKey")
|
|
if !ok {
|
|
return "", false
|
|
}
|
|
apiKey := contextStringValue(rawAPIKey)
|
|
if apiKey == "" {
|
|
return "", false
|
|
}
|
|
return apiKey, true
|
|
}
|
|
|
|
func accessMetadataSource(raw any) string {
|
|
switch v := raw.(type) {
|
|
case map[string]string:
|
|
return strings.TrimSpace(v["source"])
|
|
case map[string]any:
|
|
return contextStringValue(v["source"])
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func contextStringValue(raw any) string {
|
|
switch v := raw.(type) {
|
|
case string:
|
|
return strings.TrimSpace(v)
|
|
case []byte:
|
|
return strings.TrimSpace(string(v))
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func homeExecutionSessionIDFromMetadata(meta map[string]any) string {
|
|
if len(meta) == 0 {
|
|
return ""
|
|
}
|
|
raw, ok := meta[cliproxyexecutor.ExecutionSessionMetadataKey]
|
|
if !ok || raw == nil {
|
|
return ""
|
|
}
|
|
switch value := raw.(type) {
|
|
case string:
|
|
return strings.TrimSpace(value)
|
|
case []byte:
|
|
return strings.TrimSpace(string(value))
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func (m *Manager) clearHomeRuntimeAuths() {
|
|
if m == nil {
|
|
return
|
|
}
|
|
m.mu.Lock()
|
|
m.clearHomeRuntimeAuthsLocked()
|
|
m.mu.Unlock()
|
|
}
|
|
|
|
func (m *Manager) clearHomeRuntimeAuthsLocked() {
|
|
if m == nil {
|
|
return
|
|
}
|
|
m.homeRuntimeAuths = make(map[string]map[string]*Auth)
|
|
}
|
|
|
|
func (m *Manager) clearHomeRuntimeAuthsForSessionLocked(sessionID string) {
|
|
sessionID = strings.TrimSpace(sessionID)
|
|
if m == nil || sessionID == "" {
|
|
return
|
|
}
|
|
delete(m.homeRuntimeAuths, sessionID)
|
|
}
|
|
|
|
func (m *Manager) rememberHomeRuntimeAuth(sessionID string, auth *Auth) {
|
|
sessionID = strings.TrimSpace(sessionID)
|
|
authID := ""
|
|
if auth != nil {
|
|
authID = strings.TrimSpace(auth.ID)
|
|
}
|
|
if m == nil || auth == nil || sessionID == "" || authID == "" || !authWebsocketsEnabled(auth) {
|
|
return
|
|
}
|
|
m.mu.Lock()
|
|
if m.homeRuntimeAuths == nil {
|
|
m.homeRuntimeAuths = make(map[string]map[string]*Auth)
|
|
}
|
|
sessionAuths := m.homeRuntimeAuths[sessionID]
|
|
if sessionAuths == nil {
|
|
sessionAuths = make(map[string]*Auth)
|
|
m.homeRuntimeAuths[sessionID] = sessionAuths
|
|
}
|
|
sessionAuths[authID] = auth.Clone()
|
|
m.mu.Unlock()
|
|
}
|
|
|
|
func (m *Manager) homeRuntimeAuthByID(sessionID string, authID string) (*Auth, ProviderExecutor, string, bool) {
|
|
sessionID = strings.TrimSpace(sessionID)
|
|
authID = strings.TrimSpace(authID)
|
|
if m == nil || sessionID == "" || authID == "" {
|
|
return nil, nil, "", false
|
|
}
|
|
m.mu.RLock()
|
|
sessionAuths := m.homeRuntimeAuths[sessionID]
|
|
auth := sessionAuths[authID]
|
|
m.mu.RUnlock()
|
|
if auth == nil || !authWebsocketsEnabled(auth) {
|
|
return nil, nil, "", false
|
|
}
|
|
providerKey := executorKeyFromAuth(auth)
|
|
if providerKey == "" {
|
|
return nil, nil, "", false
|
|
}
|
|
executor, ok := m.Executor(providerKey)
|
|
if !ok && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" {
|
|
executor, ok = m.Executor("openai-compatibility")
|
|
if ok {
|
|
providerKey = "openai-compatibility"
|
|
}
|
|
}
|
|
if !ok {
|
|
return nil, nil, "", false
|
|
}
|
|
return auth.Clone(), executor, providerKey, true
|
|
}
|
|
|
|
func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) {
|
|
if m == nil {
|
|
return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"}
|
|
}
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
executionSessionID := homeExecutionSessionIDFromMetadata(opts.Metadata)
|
|
count := homeAuthCountFromMetadata(opts.Metadata)
|
|
if cliproxyexecutor.DownstreamWebsocket(ctx) && executionSessionID != "" && count <= 1 {
|
|
if pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata); pinnedAuthID != "" {
|
|
_, alreadyTried := tried[pinnedAuthID]
|
|
if !alreadyTried {
|
|
if auth, executor, providerKey, ok := m.homeRuntimeAuthByID(executionSessionID, pinnedAuthID); ok {
|
|
return auth, executor, providerKey, nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
client := currentHomeDispatcher()
|
|
if client == nil || !client.HeartbeatOK() {
|
|
return nil, nil, "", &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable}
|
|
}
|
|
|
|
requestedModel := requestedModelFromMetadata(opts.Metadata, model)
|
|
sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata)
|
|
dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers)
|
|
|
|
raw, err := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, count)
|
|
if err != nil {
|
|
if errors.Is(err, home.ErrAuthNotFound) {
|
|
return nil, nil, "", &Error{Code: "auth_not_found", Message: err.Error(), HTTPStatus: http.StatusServiceUnavailable}
|
|
}
|
|
return nil, nil, "", &Error{Code: "home_unavailable", Message: err.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable}
|
|
}
|
|
|
|
var env homeErrorEnvelope
|
|
if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal == nil && env.Error != nil {
|
|
code := strings.TrimSpace(env.Error.Type)
|
|
if code == "" {
|
|
code = strings.TrimSpace(env.Error.Code)
|
|
}
|
|
msg := strings.TrimSpace(env.Error.Message)
|
|
if msg == "" {
|
|
msg = "home returned error"
|
|
}
|
|
status := http.StatusBadGateway
|
|
switch strings.ToLower(code) {
|
|
case "model_not_found":
|
|
status = http.StatusNotFound
|
|
case "authentication_error", "unauthorized", "no_credentials", "invalid_credential":
|
|
status = http.StatusUnauthorized
|
|
}
|
|
return nil, nil, "", &Error{Code: code, Message: msg, HTTPStatus: status}
|
|
}
|
|
|
|
var dispatch homeAuthDispatchResponse
|
|
if errUnmarshal := json.Unmarshal(raw, &dispatch); errUnmarshal != nil {
|
|
return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway}
|
|
}
|
|
setHomeUserAPIKeyOnGinContext(ctx, dispatch.UserAPIKey)
|
|
auth := dispatch.Auth
|
|
if strings.TrimSpace(auth.ID) == "" {
|
|
// Backward compatibility: older home instances returned the auth directly.
|
|
if errUnmarshal := json.Unmarshal(raw, &auth); errUnmarshal != nil {
|
|
return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway}
|
|
}
|
|
}
|
|
if upstreamModel := strings.TrimSpace(dispatch.Model); upstreamModel != "" {
|
|
if auth.Attributes == nil {
|
|
auth.Attributes = make(map[string]string, 1)
|
|
}
|
|
auth.Attributes[homeUpstreamModelAttributeKey] = upstreamModel
|
|
}
|
|
if strings.TrimSpace(auth.ID) == "" {
|
|
return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned auth without id", HTTPStatus: http.StatusBadGateway}
|
|
}
|
|
if homeAuthAlreadyTried(tried, auth.ID) {
|
|
return nil, nil, "", repeatedHomeAuthError()
|
|
}
|
|
providerKey := executorKeyFromAuth(&auth)
|
|
if providerKey == "" {
|
|
return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned auth without provider", HTTPStatus: http.StatusBadGateway}
|
|
}
|
|
|
|
homeAuthIndex := strings.TrimSpace(dispatch.AuthIndex)
|
|
if homeAuthIndex != "" {
|
|
auth.Index = homeAuthIndex
|
|
auth.indexAssigned = true
|
|
} else {
|
|
auth.EnsureIndex()
|
|
}
|
|
|
|
executor, ok := m.Executor(providerKey)
|
|
if !ok && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" {
|
|
executor, ok = m.Executor("openai-compatibility")
|
|
if ok {
|
|
providerKey = "openai-compatibility"
|
|
}
|
|
}
|
|
if !ok {
|
|
return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered", HTTPStatus: http.StatusBadGateway}
|
|
}
|
|
|
|
authCopy := auth.Clone()
|
|
if cliproxyexecutor.DownstreamWebsocket(ctx) && executionSessionID != "" && authWebsocketsEnabled(authCopy) {
|
|
m.rememberHomeRuntimeAuth(executionSessionID, authCopy)
|
|
}
|
|
return authCopy, executor, providerKey, nil
|
|
}
|
|
|
|
func requestedModelFromMetadata(metadata map[string]any, fallback string) string {
|
|
if metadata != nil {
|
|
if v, ok := metadata[cliproxyexecutor.RequestedModelMetadataKey]; ok {
|
|
switch typed := v.(type) {
|
|
case string:
|
|
if trimmed := strings.TrimSpace(typed); trimmed != "" {
|
|
return trimmed
|
|
}
|
|
case []byte:
|
|
if trimmed := strings.TrimSpace(string(typed)); trimmed != "" {
|
|
return trimmed
|
|
}
|
|
}
|
|
}
|
|
}
|
|
fallback = strings.TrimSpace(fallback)
|
|
if fallback == "" {
|
|
return "unknown"
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, routeModel string, opts cliproxyexecutor.Options) ([]creditsCandidateEntry, error) {
|
|
if m == nil {
|
|
return nil, nil
|
|
}
|
|
pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata)
|
|
var candidates []creditsCandidateEntry
|
|
m.mu.RLock()
|
|
for _, auth := range m.auths {
|
|
if auth == nil || auth.Disabled || auth.Status == StatusDisabled {
|
|
continue
|
|
}
|
|
if pinnedAuthID != "" && auth.ID != pinnedAuthID {
|
|
continue
|
|
}
|
|
if !strings.EqualFold(strings.TrimSpace(auth.Provider), "antigravity") {
|
|
continue
|
|
}
|
|
if !strings.Contains(strings.ToLower(strings.TrimSpace(routeModel)), "claude") {
|
|
continue
|
|
}
|
|
providerKey := executorKeyFromAuth(auth)
|
|
executor, ok := m.executors[providerKey]
|
|
if !ok {
|
|
continue
|
|
}
|
|
candidates = append(candidates, creditsCandidateEntry{
|
|
auth: auth.Clone(),
|
|
executor: executor,
|
|
provider: providerKey,
|
|
})
|
|
}
|
|
m.mu.RUnlock()
|
|
|
|
var known []creditsCandidateEntry
|
|
var unknown []creditsCandidateEntry
|
|
for _, candidate := range candidates {
|
|
hint, okHint, errHint := GetAntigravityCreditsHintRequired(ctx, candidate.auth.ID)
|
|
if errHint != nil {
|
|
return nil, antigravityCreditsKVUnavailableError(errHint)
|
|
}
|
|
if okHint && hint.Known {
|
|
if !hint.Available {
|
|
continue
|
|
}
|
|
known = append(known, candidate)
|
|
continue
|
|
}
|
|
unknown = append(unknown, candidate)
|
|
}
|
|
sort.Slice(known, func(i, j int) bool {
|
|
return known[i].auth.ID < known[j].auth.ID
|
|
})
|
|
sort.Slice(unknown, func(i, j int) bool {
|
|
return unknown[i].auth.ID < unknown[j].auth.ID
|
|
})
|
|
return append(known, unknown...), nil
|
|
}
|
|
|
|
type creditsCandidateEntry struct {
|
|
auth *Auth
|
|
executor ProviderExecutor
|
|
provider string
|
|
}
|
|
|
|
func hasAntigravityProvider(providers []string) bool {
|
|
for _, p := range providers {
|
|
if strings.EqualFold(strings.TrimSpace(p), "antigravity") {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, providers []string) bool {
|
|
status := statusCodeFromError(lastErr)
|
|
log.WithFields(log.Fields{
|
|
"lastErr": errorString(lastErr),
|
|
"status": status,
|
|
"providers": providers,
|
|
}).Debug("shouldAttemptAntigravityCreditsFallback")
|
|
if m == nil || lastErr == nil {
|
|
return false
|
|
}
|
|
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
|
if cfg == nil || !cfg.QuotaExceeded.AntigravityCredits {
|
|
return false
|
|
}
|
|
switch status {
|
|
case http.StatusTooManyRequests, http.StatusServiceUnavailable:
|
|
return true
|
|
case 0:
|
|
var authErr *Error
|
|
if errors.As(lastErr, &authErr) && authErr != nil {
|
|
return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" || authErr.Code == "model_cooldown"
|
|
}
|
|
var cooldownErr *modelCooldownError
|
|
if errors.As(lastErr, &cooldownErr) {
|
|
return true
|
|
}
|
|
return false
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, bool, error) {
|
|
routeModel := req.Model
|
|
candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts)
|
|
if errCandidates != nil {
|
|
return cliproxyexecutor.Response{}, false, errCandidates
|
|
}
|
|
for _, c := range candidates {
|
|
if ctx.Err() != nil {
|
|
return cliproxyexecutor.Response{}, false, nil
|
|
}
|
|
creditsCtx := WithAntigravityCredits(ctx)
|
|
if rt := m.roundTripperFor(c.auth); rt != nil {
|
|
creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt)
|
|
creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt)
|
|
}
|
|
creditsOpts := ensureRequestedModelMetadata(opts, routeModel)
|
|
creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel)
|
|
preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth)
|
|
if errPrepare != nil {
|
|
continue
|
|
}
|
|
c.auth = preparedAuth
|
|
publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID)
|
|
models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel)
|
|
if len(models) == 0 {
|
|
continue
|
|
}
|
|
for _, upstreamModel := range models {
|
|
resultModel := m.stateModelForExecution(c.auth, routeModel, upstreamModel, pooled)
|
|
execReq := req
|
|
execReq.Model = upstreamModel
|
|
resp, errExec := c.executor.Execute(creditsCtx, c.auth, execReq, creditsOpts)
|
|
result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, Success: errExec == nil}
|
|
if errExec != nil {
|
|
result.Error = &Error{Message: errExec.Error()}
|
|
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errExec); ok && se != nil {
|
|
result.Error.HTTPStatus = se.StatusCode()
|
|
}
|
|
if ra := retryAfterFromError(errExec); ra != nil {
|
|
result.RetryAfter = ra
|
|
}
|
|
m.MarkResult(creditsCtx, result)
|
|
continue
|
|
}
|
|
m.MarkResult(creditsCtx, result)
|
|
rewriteForceMappedResponse(&resp, aliasResult)
|
|
return resp, true, nil
|
|
}
|
|
}
|
|
return cliproxyexecutor.Response{}, false, nil
|
|
}
|
|
|
|
func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, bool, error) {
|
|
routeModel := req.Model
|
|
candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts)
|
|
if errCandidates != nil {
|
|
return nil, false, errCandidates
|
|
}
|
|
for _, c := range candidates {
|
|
if ctx.Err() != nil {
|
|
return nil, false, nil
|
|
}
|
|
creditsCtx := WithAntigravityCredits(ctx)
|
|
if rt := m.roundTripperFor(c.auth); rt != nil {
|
|
creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt)
|
|
creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt)
|
|
}
|
|
creditsOpts := ensureRequestedModelMetadata(opts, routeModel)
|
|
preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth)
|
|
if errPrepare != nil {
|
|
continue
|
|
}
|
|
c.auth = preparedAuth
|
|
publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID)
|
|
models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel)
|
|
if len(models) == 0 {
|
|
continue
|
|
}
|
|
result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, "", models, pooled, aliasResult)
|
|
if errStream != nil {
|
|
continue
|
|
}
|
|
return result, true, nil
|
|
}
|
|
return nil, false, nil
|
|
}
|
|
|
|
func antigravityCreditsKVUnavailableError(cause error) error {
|
|
if cause == nil {
|
|
return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable", HTTPStatus: http.StatusServiceUnavailable}
|
|
}
|
|
return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable: " + cause.Error(), HTTPStatus: http.StatusServiceUnavailable}
|
|
}
|
|
|
|
func (m *Manager) persist(ctx context.Context, auth *Auth) error {
|
|
if m.store == nil || auth == nil {
|
|
return nil
|
|
}
|
|
if shouldSkipPersist(ctx) {
|
|
return nil
|
|
}
|
|
if IsConfigAPIKeyAuth(auth) {
|
|
return nil
|
|
}
|
|
if auth.Attributes != nil {
|
|
if v := strings.ToLower(strings.TrimSpace(auth.Attributes["runtime_only"])); v == "true" {
|
|
return nil
|
|
}
|
|
}
|
|
if IsPluginVirtualAuth(auth) {
|
|
return nil
|
|
}
|
|
// Skip persistence when metadata is absent (e.g., runtime-only auths).
|
|
if auth.Metadata == nil {
|
|
return nil
|
|
}
|
|
_, err := m.store.Save(ctx, auth)
|
|
return err
|
|
}
|
|
|
|
// StartAutoRefresh launches a background loop that evaluates auth freshness
|
|
// every few seconds and triggers refresh operations when required.
|
|
// Only one loop is kept alive; starting a new one cancels the previous run.
|
|
func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duration) {
|
|
if interval <= 0 {
|
|
interval = refreshCheckInterval
|
|
}
|
|
|
|
m.mu.Lock()
|
|
cancelPrev := m.refreshCancel
|
|
m.refreshCancel = nil
|
|
m.refreshLoop = nil
|
|
m.mu.Unlock()
|
|
if cancelPrev != nil {
|
|
cancelPrev()
|
|
}
|
|
|
|
ctx, cancelCtx := context.WithCancel(parent)
|
|
workers := refreshMaxConcurrency
|
|
if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil && cfg.AuthAutoRefreshWorkers > 0 {
|
|
workers = cfg.AuthAutoRefreshWorkers
|
|
}
|
|
loop := newAuthAutoRefreshLoop(m, interval, workers)
|
|
|
|
m.mu.Lock()
|
|
m.refreshCancel = cancelCtx
|
|
m.refreshLoop = loop
|
|
m.mu.Unlock()
|
|
|
|
loop.rebuild(time.Now())
|
|
go loop.run(ctx)
|
|
}
|
|
|
|
// StopAutoRefresh cancels the background refresh loop, if running.
|
|
// It also stops the selector if it implements StoppableSelector.
|
|
func (m *Manager) StopAutoRefresh() {
|
|
m.mu.Lock()
|
|
cancel := m.refreshCancel
|
|
m.refreshCancel = nil
|
|
m.refreshLoop = nil
|
|
m.mu.Unlock()
|
|
if cancel != nil {
|
|
cancel()
|
|
}
|
|
// Stop selector if it implements StoppableSelector (e.g., SessionAffinitySelector)
|
|
if stoppable, ok := m.selector.(StoppableSelector); ok {
|
|
stoppable.Stop()
|
|
}
|
|
}
|
|
|
|
func (m *Manager) queueRefreshReschedule(authID string) {
|
|
if m == nil || authID == "" {
|
|
return
|
|
}
|
|
m.mu.RLock()
|
|
loop := m.refreshLoop
|
|
m.mu.RUnlock()
|
|
if loop == nil {
|
|
return
|
|
}
|
|
loop.queueReschedule(authID)
|
|
}
|
|
|
|
func (m *Manager) queueRefreshUnschedule(authID string) {
|
|
if m == nil || authID == "" {
|
|
return
|
|
}
|
|
m.mu.RLock()
|
|
loop := m.refreshLoop
|
|
m.mu.RUnlock()
|
|
if loop == nil {
|
|
return
|
|
}
|
|
loop.remove(authID)
|
|
}
|
|
|
|
func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool {
|
|
if a == nil {
|
|
return false
|
|
}
|
|
if hasUnauthorizedAuthFailure(a) {
|
|
return false
|
|
}
|
|
if !a.NextRefreshAfter.IsZero() && now.Before(a.NextRefreshAfter) {
|
|
return false
|
|
}
|
|
if evaluator, ok := a.Runtime.(RefreshEvaluator); ok && evaluator != nil {
|
|
return evaluator.ShouldRefresh(now, a)
|
|
}
|
|
|
|
lastRefresh := a.LastRefreshedAt
|
|
if lastRefresh.IsZero() {
|
|
if ts, ok := authLastRefreshTimestamp(a); ok {
|
|
lastRefresh = ts
|
|
}
|
|
}
|
|
|
|
expiry, hasExpiry := a.ExpirationTime()
|
|
|
|
if interval := authPreferredInterval(a); interval > 0 {
|
|
if hasExpiry && !expiry.IsZero() {
|
|
if !expiry.After(now) {
|
|
return true
|
|
}
|
|
if expiry.Sub(now) <= interval {
|
|
return true
|
|
}
|
|
}
|
|
if lastRefresh.IsZero() {
|
|
return true
|
|
}
|
|
return now.Sub(lastRefresh) >= interval
|
|
}
|
|
|
|
provider := strings.ToLower(a.Provider)
|
|
lead := ProviderRefreshLead(provider, a.Runtime)
|
|
if lead == nil {
|
|
return false
|
|
}
|
|
if *lead <= 0 {
|
|
if hasExpiry && !expiry.IsZero() {
|
|
return now.After(expiry)
|
|
}
|
|
return false
|
|
}
|
|
if hasExpiry && !expiry.IsZero() {
|
|
return time.Until(expiry) <= *lead
|
|
}
|
|
if !lastRefresh.IsZero() {
|
|
return now.Sub(lastRefresh) >= *lead
|
|
}
|
|
return true
|
|
}
|
|
|
|
func authPreferredInterval(a *Auth) time.Duration {
|
|
if a == nil {
|
|
return 0
|
|
}
|
|
if d := durationFromMetadata(a.Metadata, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 {
|
|
return d
|
|
}
|
|
if d := durationFromAttributes(a.Attributes, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 {
|
|
return d
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func durationFromMetadata(meta map[string]any, keys ...string) time.Duration {
|
|
if len(meta) == 0 {
|
|
return 0
|
|
}
|
|
for _, key := range keys {
|
|
if val, ok := meta[key]; ok {
|
|
if dur := parseDurationValue(val); dur > 0 {
|
|
return dur
|
|
}
|
|
}
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func durationFromAttributes(attrs map[string]string, keys ...string) time.Duration {
|
|
if len(attrs) == 0 {
|
|
return 0
|
|
}
|
|
for _, key := range keys {
|
|
if val, ok := attrs[key]; ok {
|
|
if dur := parseDurationString(val); dur > 0 {
|
|
return dur
|
|
}
|
|
}
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func parseDurationValue(val any) time.Duration {
|
|
switch v := val.(type) {
|
|
case time.Duration:
|
|
if v <= 0 {
|
|
return 0
|
|
}
|
|
return v
|
|
case int:
|
|
if v <= 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(v) * time.Second
|
|
case int32:
|
|
if v <= 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(v) * time.Second
|
|
case int64:
|
|
if v <= 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(v) * time.Second
|
|
case uint:
|
|
if v == 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(v) * time.Second
|
|
case uint32:
|
|
if v == 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(v) * time.Second
|
|
case uint64:
|
|
if v == 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(v) * time.Second
|
|
case float32:
|
|
if v <= 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(float64(v) * float64(time.Second))
|
|
case float64:
|
|
if v <= 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(v * float64(time.Second))
|
|
case json.Number:
|
|
if i, err := v.Int64(); err == nil {
|
|
if i <= 0 {
|
|
return 0
|
|
}
|
|
return time.Duration(i) * time.Second
|
|
}
|
|
if f, err := v.Float64(); err == nil && f > 0 {
|
|
return time.Duration(f * float64(time.Second))
|
|
}
|
|
case string:
|
|
return parseDurationString(v)
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func parseDurationString(raw string) time.Duration {
|
|
s := strings.TrimSpace(raw)
|
|
if s == "" {
|
|
return 0
|
|
}
|
|
if dur, err := time.ParseDuration(s); err == nil && dur > 0 {
|
|
return dur
|
|
}
|
|
if secs, err := strconv.ParseFloat(s, 64); err == nil && secs > 0 {
|
|
return time.Duration(secs * float64(time.Second))
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func authLastRefreshTimestamp(a *Auth) (time.Time, bool) {
|
|
if a == nil {
|
|
return time.Time{}, false
|
|
}
|
|
if a.Metadata != nil {
|
|
if ts, ok := lookupMetadataTime(a.Metadata, "last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"); ok {
|
|
return ts, true
|
|
}
|
|
}
|
|
if a.Attributes != nil {
|
|
for _, key := range []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"} {
|
|
if val := strings.TrimSpace(a.Attributes[key]); val != "" {
|
|
if ts, ok := parseTimeValue(val); ok {
|
|
return ts, true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return time.Time{}, false
|
|
}
|
|
|
|
func lookupMetadataTime(meta map[string]any, keys ...string) (time.Time, bool) {
|
|
for _, key := range keys {
|
|
if val, ok := meta[key]; ok {
|
|
if ts, ok1 := parseTimeValue(val); ok1 {
|
|
return ts, true
|
|
}
|
|
}
|
|
}
|
|
return time.Time{}, false
|
|
}
|
|
|
|
func (m *Manager) markRefreshPending(id string, now time.Time) bool {
|
|
m.mu.Lock()
|
|
auth, ok := m.auths[id]
|
|
if !ok || auth == nil {
|
|
m.mu.Unlock()
|
|
return false
|
|
}
|
|
if !auth.NextRefreshAfter.IsZero() && now.Before(auth.NextRefreshAfter) {
|
|
m.mu.Unlock()
|
|
return false
|
|
}
|
|
auth.NextRefreshAfter = now.Add(refreshPendingBackoff)
|
|
m.auths[id] = auth
|
|
m.mu.Unlock()
|
|
|
|
m.queueRefreshReschedule(id)
|
|
return true
|
|
}
|
|
|
|
type authRefreshLock struct {
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func authAccessToken(auth *Auth) string {
|
|
if token := authMetadataString(auth, "access_token"); token != "" {
|
|
return token
|
|
}
|
|
return authMetadataString(auth, "accessToken")
|
|
}
|
|
|
|
func authHasRefreshCredential(auth *Auth) bool {
|
|
if authMetadataString(auth, "refresh_token") != "" {
|
|
return true
|
|
}
|
|
return authMetadataString(auth, "refreshToken") != ""
|
|
}
|
|
|
|
func clearUnauthorizedModelStates(auth *Auth, now time.Time) []string {
|
|
if auth == nil || len(auth.ModelStates) == 0 {
|
|
return nil
|
|
}
|
|
var resumed []string
|
|
for model, state := range auth.ModelStates {
|
|
if state == nil || state.LastError == nil {
|
|
continue
|
|
}
|
|
if state.LastError.StatusCode() != http.StatusUnauthorized && !strings.EqualFold(state.LastError.Code, "unauthorized") {
|
|
continue
|
|
}
|
|
resetModelState(state, now)
|
|
resumed = append(resumed, model)
|
|
}
|
|
if len(resumed) > 0 {
|
|
updateAggregatedAvailability(auth, now)
|
|
}
|
|
return resumed
|
|
}
|
|
|
|
// tryRefreshAfterUnauthorized refreshes OAuth credentials once after a 401 so the
|
|
// current auth can be retried before fallback/suspend.
|
|
func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, execErr error, alreadyTried bool) (*Auth, bool) {
|
|
if m == nil || auth == nil || alreadyTried || execErr == nil {
|
|
return auth, false
|
|
}
|
|
if !isUnauthorizedError(execErr) || !authHasRefreshCredential(auth) {
|
|
return auth, false
|
|
}
|
|
log.Debugf("unauthorized response for %s (%s), refreshing credentials before fallback", auth.Provider, auth.ID)
|
|
refreshed, errRefresh := m.refreshAuthForRequest(ctx, auth.ID, authAccessToken(auth))
|
|
if errRefresh != nil || refreshed == nil {
|
|
log.Debugf("credential refresh before fallback failed for %s (%s): %v", auth.Provider, auth.ID, errRefresh)
|
|
return auth, false
|
|
}
|
|
return refreshed, true
|
|
}
|
|
|
|
func (m *Manager) refreshAuth(ctx context.Context, id string) {
|
|
_, _ = m.refreshAuthForRequest(ctx, id, "")
|
|
}
|
|
|
|
// refreshAuthForRequest performs a synchronous credential refresh for the given auth.
|
|
// failedAccessToken lets concurrent callers reuse a refresh that already replaced the
|
|
// access token that produced the unauthorized response.
|
|
func (m *Manager) refreshAuthForRequest(ctx context.Context, id, failedAccessToken string) (*Auth, error) {
|
|
if m == nil {
|
|
return nil, errors.New("auth manager is nil")
|
|
}
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
id = strings.TrimSpace(id)
|
|
if id == "" {
|
|
return nil, errors.New("auth id is empty")
|
|
}
|
|
|
|
lockValue, _ := m.refreshLocks.LoadOrStore(id, &authRefreshLock{})
|
|
lock, _ := lockValue.(*authRefreshLock)
|
|
if lock == nil {
|
|
lock = &authRefreshLock{}
|
|
m.refreshLocks.Store(id, lock)
|
|
}
|
|
lock.mu.Lock()
|
|
defer lock.mu.Unlock()
|
|
|
|
m.mu.RLock()
|
|
auth := m.auths[id]
|
|
var exec ProviderExecutor
|
|
if auth != nil {
|
|
exec = m.executors[auth.Provider]
|
|
}
|
|
m.mu.RUnlock()
|
|
if auth == nil || exec == nil {
|
|
return nil, errors.New("auth or executor not found")
|
|
}
|
|
|
|
// Another request may already have refreshed this credential.
|
|
if failedAccessToken != "" {
|
|
if currentToken := authAccessToken(auth); currentToken != "" && currentToken != failedAccessToken {
|
|
return auth.Clone(), nil
|
|
}
|
|
}
|
|
|
|
cloned := auth.Clone()
|
|
updated, err := exec.Refresh(ctx, cloned)
|
|
if err != nil && errors.Is(err, context.Canceled) {
|
|
log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID)
|
|
return nil, err
|
|
}
|
|
log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err)
|
|
now := time.Now()
|
|
if err != nil {
|
|
unauthorized := isUnauthorizedError(err)
|
|
shouldReschedule := false
|
|
m.mu.Lock()
|
|
if current := m.auths[id]; current != nil {
|
|
current.LastError = refreshErrorFromError(err)
|
|
if unauthorized {
|
|
current.NextRefreshAfter = time.Time{}
|
|
current.Unavailable = true
|
|
current.Status = StatusError
|
|
current.StatusMessage = "unauthorized"
|
|
} else {
|
|
current.NextRefreshAfter = now.Add(refreshFailureBackoff)
|
|
}
|
|
m.auths[id] = current
|
|
shouldReschedule = true
|
|
if m.scheduler != nil {
|
|
m.scheduler.upsertAuth(current.Clone())
|
|
}
|
|
}
|
|
m.mu.Unlock()
|
|
if shouldReschedule {
|
|
m.queueRefreshReschedule(id)
|
|
}
|
|
return nil, err
|
|
}
|
|
if updated == nil {
|
|
updated = cloned
|
|
}
|
|
// Preserve runtime created by the executor during Refresh.
|
|
// If executor didn't set one, fall back to the previous runtime.
|
|
if updated.Runtime == nil {
|
|
updated.Runtime = auth.Runtime
|
|
}
|
|
updated.LastRefreshedAt = now
|
|
updated.NextRefreshAfter = time.Time{}
|
|
updated.LastError = nil
|
|
updated.StatusMessage = ""
|
|
updated.Unavailable = false
|
|
if updated.Status == StatusError {
|
|
updated.Status = StatusActive
|
|
}
|
|
updated.UpdatedAt = now
|
|
modelsToResume := clearUnauthorizedModelStates(updated, now)
|
|
if m.shouldRefresh(updated, now) {
|
|
updated.NextRefreshAfter = now.Add(refreshIneffectiveBackoff)
|
|
}
|
|
saved, errUpdate := m.Update(ctx, updated)
|
|
for _, model := range modelsToResume {
|
|
registry.GetGlobalRegistry().ResumeClientModel(id, model)
|
|
}
|
|
if errUpdate != nil {
|
|
log.Debugf("persist refreshed auth %s (%s) failed: %v", auth.Provider, auth.ID, errUpdate)
|
|
}
|
|
if saved != nil {
|
|
return saved, nil
|
|
}
|
|
return updated.Clone(), nil
|
|
}
|
|
|
|
func (m *Manager) executorFor(provider string) ProviderExecutor {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
return m.executors[provider]
|
|
}
|
|
|
|
// roundTripperContextKey is an unexported context key type to avoid collisions.
|
|
type roundTripperContextKey struct{}
|
|
|
|
// roundTripperFor retrieves an HTTP RoundTripper for the given auth if a provider is registered.
|
|
func (m *Manager) roundTripperFor(auth *Auth) http.RoundTripper {
|
|
m.mu.RLock()
|
|
p := m.rtProvider
|
|
m.mu.RUnlock()
|
|
if p == nil || auth == nil {
|
|
return nil
|
|
}
|
|
return p.RoundTripperFor(auth)
|
|
}
|
|
|
|
// RoundTripperProvider defines a minimal provider of per-auth HTTP transports.
|
|
type RoundTripperProvider interface {
|
|
RoundTripperFor(auth *Auth) http.RoundTripper
|
|
}
|
|
|
|
// RequestPreparer is an optional interface that provider executors can implement
|
|
// to mutate outbound HTTP requests with provider credentials.
|
|
type RequestPreparer interface {
|
|
PrepareRequest(req *http.Request, auth *Auth) error
|
|
}
|
|
|
|
func executorKeyFromAuth(auth *Auth) string {
|
|
if auth == nil {
|
|
return ""
|
|
}
|
|
if auth.Attributes != nil {
|
|
providerKey := strings.TrimSpace(auth.Attributes["provider_key"])
|
|
compatName := strings.TrimSpace(auth.Attributes["compat_name"])
|
|
if compatName != "" {
|
|
if providerKey == "" {
|
|
providerKey = compatName
|
|
}
|
|
return util.OpenAICompatibleProviderKey(providerKey)
|
|
}
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
|
|
providerKey := strings.TrimSpace(auth.Label)
|
|
if providerKey == "" {
|
|
providerKey = "openai-compatibility"
|
|
}
|
|
return util.OpenAICompatibleProviderKey(providerKey)
|
|
}
|
|
return strings.ToLower(strings.TrimSpace(auth.Provider))
|
|
}
|
|
|
|
// logEntryWithRequestID returns a logrus entry with request_id field if available in context.
|
|
func logEntryWithRequestID(ctx context.Context) *log.Entry {
|
|
if ctx == nil {
|
|
return log.NewEntry(log.StandardLogger())
|
|
}
|
|
if reqID := logging.GetRequestID(ctx); reqID != "" {
|
|
return log.WithField("request_id", reqID)
|
|
}
|
|
return log.NewEntry(log.StandardLogger())
|
|
}
|
|
|
|
func debugLogAuthSelection(entry *log.Entry, auth *Auth, provider string, model string) {
|
|
if !log.IsLevelEnabled(log.DebugLevel) {
|
|
return
|
|
}
|
|
if entry == nil || auth == nil {
|
|
return
|
|
}
|
|
accountType, accountInfo := auth.AccountInfo()
|
|
proxyInfo := auth.ProxyInfo()
|
|
suffix := ""
|
|
if proxyInfo != "" {
|
|
suffix = " " + proxyInfo
|
|
}
|
|
switch accountType {
|
|
case "api_key":
|
|
entry.Debugf("Use API key %s for model %s%s", util.HideAPIKey(accountInfo), model, suffix)
|
|
case "oauth":
|
|
ident := formatOauthIdentity(auth, provider, accountInfo)
|
|
entry.Debugf("Use OAuth %s for model %s%s", ident, model, suffix)
|
|
}
|
|
}
|
|
|
|
func formatOauthIdentity(auth *Auth, provider string, accountInfo string) string {
|
|
if auth == nil {
|
|
return ""
|
|
}
|
|
// Prefer the auth's provider when available.
|
|
providerName := strings.TrimSpace(auth.Provider)
|
|
if providerName == "" {
|
|
providerName = strings.TrimSpace(provider)
|
|
}
|
|
// Only log the basename to avoid leaking host paths.
|
|
// FileName may be unset for some auth backends; fall back to ID.
|
|
authFile := strings.TrimSpace(auth.FileName)
|
|
if authFile == "" {
|
|
authFile = strings.TrimSpace(auth.ID)
|
|
}
|
|
if authFile != "" {
|
|
authFile = filepath.Base(authFile)
|
|
}
|
|
parts := make([]string, 0, 3)
|
|
if providerName != "" {
|
|
parts = append(parts, "provider="+providerName)
|
|
}
|
|
if authFile != "" {
|
|
parts = append(parts, "auth_file="+authFile)
|
|
}
|
|
if len(parts) == 0 {
|
|
return accountInfo
|
|
}
|
|
return strings.Join(parts, " ")
|
|
}
|
|
|
|
// InjectCredentials delegates per-provider HTTP request preparation when supported.
|
|
// If the registered executor for the auth provider implements RequestPreparer,
|
|
// it will be invoked to modify the request (e.g., add headers).
|
|
func (m *Manager) InjectCredentials(req *http.Request, authID string) error {
|
|
if req == nil || authID == "" {
|
|
return nil
|
|
}
|
|
m.mu.RLock()
|
|
a := m.auths[authID]
|
|
var exec ProviderExecutor
|
|
if a != nil {
|
|
exec = m.executors[executorKeyFromAuth(a)]
|
|
}
|
|
m.mu.RUnlock()
|
|
if a == nil || exec == nil {
|
|
return nil
|
|
}
|
|
if p, ok := exec.(RequestPreparer); ok && p != nil {
|
|
return p.PrepareRequest(req, a)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// PrepareHttpRequest injects provider credentials into the supplied HTTP request.
|
|
func (m *Manager) PrepareHttpRequest(ctx context.Context, auth *Auth, req *http.Request) error {
|
|
if m == nil {
|
|
return &Error{Code: "provider_not_found", Message: "manager is nil"}
|
|
}
|
|
if auth == nil {
|
|
return &Error{Code: "auth_not_found", Message: "auth is nil"}
|
|
}
|
|
if req == nil {
|
|
return &Error{Code: "invalid_request", Message: "http request is nil"}
|
|
}
|
|
if ctx != nil {
|
|
*req = *req.WithContext(ctx)
|
|
}
|
|
providerKey := executorKeyFromAuth(auth)
|
|
if providerKey == "" {
|
|
return &Error{Code: "provider_not_found", Message: "auth provider is empty"}
|
|
}
|
|
exec := m.executorFor(providerKey)
|
|
if exec == nil {
|
|
return &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey}
|
|
}
|
|
preparer, ok := exec.(RequestPreparer)
|
|
if !ok || preparer == nil {
|
|
return &Error{Code: "not_supported", Message: "executor does not support http request preparation"}
|
|
}
|
|
return preparer.PrepareRequest(req, auth)
|
|
}
|
|
|
|
// NewHttpRequest constructs a new HTTP request and injects provider credentials into it.
|
|
func (m *Manager) NewHttpRequest(ctx context.Context, auth *Auth, method, targetURL string, body []byte, headers http.Header) (*http.Request, error) {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
method = strings.TrimSpace(method)
|
|
if method == "" {
|
|
method = http.MethodGet
|
|
}
|
|
var reader io.Reader
|
|
if body != nil {
|
|
reader = bytes.NewReader(body)
|
|
}
|
|
httpReq, err := http.NewRequestWithContext(ctx, method, targetURL, reader)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if headers != nil {
|
|
httpReq.Header = headers.Clone()
|
|
}
|
|
if errPrepare := m.PrepareHttpRequest(ctx, auth, httpReq); errPrepare != nil {
|
|
return nil, errPrepare
|
|
}
|
|
return httpReq, nil
|
|
}
|
|
|
|
// HttpRequest injects provider credentials into the supplied HTTP request and executes it.
|
|
func (m *Manager) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) {
|
|
if m == nil {
|
|
return nil, &Error{Code: "provider_not_found", Message: "manager is nil"}
|
|
}
|
|
if auth == nil {
|
|
return nil, &Error{Code: "auth_not_found", Message: "auth is nil"}
|
|
}
|
|
if req == nil {
|
|
return nil, &Error{Code: "invalid_request", Message: "http request is nil"}
|
|
}
|
|
providerKey := executorKeyFromAuth(auth)
|
|
if providerKey == "" {
|
|
return nil, &Error{Code: "provider_not_found", Message: "auth provider is empty"}
|
|
}
|
|
exec := m.executorFor(providerKey)
|
|
if exec == nil {
|
|
return nil, &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey}
|
|
}
|
|
return exec.HttpRequest(ctx, auth, req)
|
|
}
|