Add projects

backend/sdk/cliproxy/auth/conductor.go

@@ -0,0 +1,195 @@
+package auth
+
+import (
+	"context"
+	"net/http"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
+	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
+	"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
+)
+
+// 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)
+}
+
+// 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
+	// CredentialScope indicates that the failure affects the whole credential across models (e.g. Anthropic 5h/7d unified limits).
+	CredentialScope bool
+	// Error describes the failure when Success is false.
+	Error *Error
+	// Options carries execution request options (headers, metadata, etc.) for result tracking.
+	Options cliproxyexecutor.Options
+}
+
+// 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
+	pendingCooldownStateStore CooldownStateStore
+	executors                 map[string]ProviderExecutor
+	selector                  Selector
+	hook                      Hook
+	mu                        sync.RWMutex
+	selectorMu                sync.Mutex
+	configCooldownMu          sync.Mutex
+	auths                     map[string]*Auth
+	scheduler                 *authScheduler
+	// pluginScheduler runs outside m.mu before falling back to native selection.
+	pluginScheduler PluginScheduler
+	// homeRuntimeAuths retains legacy session auth lookups for non-execution callers.
+	homeRuntimeAuths map[string]map[string]*Auth
+	// homeRuntimeAuthOwners prevents a stale selection from clearing a replacement auth.
+	homeRuntimeAuthOwners map[string]map[string]*HomeDispatchSelection
+	// homeSessionSelections owns retained Home selections for websocket sessions.
+	homeSessionSelections map[string]map[homeSessionSelectionKey]*HomeDispatchSelection
+	homeSessionLocks      sync.Map
+	homeSessionAliases    homeSessionAliasCache
+	// providerOffsets tracks per-model provider rotation state for multi-provider routing.
+	providerOffsets             map[string]int
+	homeDispatchBundle          atomic.Pointer[HomeDispatchBundle]
+	homeInFlightPublisherConfig atomic.Pointer[HomeInFlightPublisherConfig]
+
+	// 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
+
+	// apiKeyModelRouting atomically publishes per-auth aliases and configured capabilities.
+	apiKeyModelRouting 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),
+		homeRuntimeAuthOwners: make(map[string]map[string]*HomeDispatchSelection),
+		homeSessionSelections: make(map[string]map[homeSessionSelectionKey]*HomeDispatchSelection),
+		providerOffsets:       make(map[string]int),
+		modelPoolOffsets:      make(map[string]int),
+	}
+	// atomic.Value requires non-nil initial value.
+	manager.runtimeConfig.Store(&internalconfig.Config{})
+	manager.apiKeyModelRouting.Store(&apiKeyModelRoutingSnapshot{config: &internalconfig.Config{}})
+	defaultInFlightConfig, errInFlightConfig := HomeInFlightPublisherConfigFromConfig(internalconfig.DefaultCredentialInFlightConfig())
+	if errInFlightConfig == nil {
+		manager.ApplyHomeInFlightPublisherConfig(defaultInFlightConfig)
+	}
+	manager.scheduler = newAuthScheduler(selector)
+	return manager
+}