rebuild branch

This commit is contained in:
Luis Pater
2025-09-25 10:32:48 +08:00
parent 3f69254f43
commit f5dc380b63
214 changed files with 39377 additions and 0 deletions

View File

@@ -0,0 +1,32 @@
package auth
// Error describes an authentication related failure in a provider agnostic format.
type Error struct {
// Code is a short machine readable identifier.
Code string `json:"code,omitempty"`
// Message is a human readable description of the failure.
Message string `json:"message"`
// Retryable indicates whether a retry might fix the issue automatically.
Retryable bool `json:"retryable"`
// HTTPStatus optionally records an HTTP-like status code for the error.
HTTPStatus int `json:"http_status,omitempty"`
}
// Error implements the error interface.
func (e *Error) Error() string {
if e == nil {
return ""
}
if e.Code == "" {
return e.Message
}
return e.Code + ": " + e.Message
}
// StatusCode implements optional status accessor for manager decision making.
func (e *Error) StatusCode() int {
if e == nil {
return 0
}
return e.HTTPStatus
}

1206
sdk/cliproxy/auth/manager.go Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,79 @@
package auth
import (
"context"
"sync"
"time"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
)
// RoundRobinSelector provides a simple provider scoped round-robin selection strategy.
type RoundRobinSelector struct {
mu sync.Mutex
cursors map[string]int
}
// Pick selects the next available auth for the provider in a round-robin manner.
func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
_ = ctx
_ = opts
if len(auths) == 0 {
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
}
if s.cursors == nil {
s.cursors = make(map[string]int)
}
available := make([]*Auth, 0, len(auths))
now := time.Now()
for i := 0; i < len(auths); i++ {
candidate := auths[i]
if isAuthBlockedForModel(candidate, model, now) {
continue
}
available = append(available, candidate)
}
if len(available) == 0 {
return nil, &Error{Code: "auth_unavailable", Message: "no auth available"}
}
key := provider + ":" + model
s.mu.Lock()
index := s.cursors[key]
if index >= 2_147_483_640 {
index = 0
}
s.cursors[key] = index + 1
s.mu.Unlock()
// log.Debugf("available: %d, index: %d, key: %d", len(available), index, index%len(available))
return available[index%len(available)], nil
}
func isAuthBlockedForModel(auth *Auth, model string, now time.Time) bool {
if auth == nil {
return true
}
if auth.Disabled || auth.Status == StatusDisabled {
return true
}
if model != "" && len(auth.ModelStates) > 0 {
if state, ok := auth.ModelStates[model]; ok && state != nil {
if state.Status == StatusDisabled {
return true
}
if state.Unavailable {
if state.NextRetryAfter.IsZero() {
return false
}
if state.NextRetryAfter.After(now) {
return true
}
}
}
}
if auth.Unavailable && auth.NextRetryAfter.After(now) {
return true
}
return false
}

View File

@@ -0,0 +1,19 @@
package auth
// Status represents the lifecycle state of an Auth entry.
type Status string
const (
// StatusUnknown means the auth state could not be determined.
StatusUnknown Status = "unknown"
// StatusActive indicates the auth is valid and ready for execution.
StatusActive Status = "active"
// StatusPending indicates the auth is waiting for an external action, such as MFA.
StatusPending Status = "pending"
// StatusRefreshing indicates the auth is undergoing a refresh flow.
StatusRefreshing Status = "refreshing"
// StatusError indicates the auth is temporarily unavailable due to errors.
StatusError Status = "error"
// StatusDisabled marks the auth as intentionally disabled.
StatusDisabled Status = "disabled"
)

View File

@@ -0,0 +1,13 @@
package auth
import "context"
// Store abstracts persistence of Auth state across restarts.
type Store interface {
// List returns all auth records stored in the backend.
List(ctx context.Context) ([]*Auth, error)
// SaveAuth persists the provided auth record, replacing any existing one with same ID.
SaveAuth(ctx context.Context, auth *Auth) error
// Delete removes the auth record identified by id.
Delete(ctx context.Context, id string) error
}

289
sdk/cliproxy/auth/types.go Normal file
View File

@@ -0,0 +1,289 @@
package auth
import (
"encoding/json"
"strconv"
"strings"
"sync"
"time"
)
// Auth encapsulates the runtime state and metadata associated with a single credential.
type Auth struct {
// ID uniquely identifies the auth record across restarts.
ID string `json:"id"`
// Provider is the upstream provider key (e.g. "gemini", "claude").
Provider string `json:"provider"`
// Label is an optional human readable label for logging.
Label string `json:"label,omitempty"`
// Status is the lifecycle status managed by the AuthManager.
Status Status `json:"status"`
// StatusMessage holds a short description for the current status.
StatusMessage string `json:"status_message,omitempty"`
// Disabled indicates the auth is intentionally disabled by operator.
Disabled bool `json:"disabled"`
// Unavailable flags transient provider unavailability (e.g. quota exceeded).
Unavailable bool `json:"unavailable"`
// ProxyURL overrides the global proxy setting for this auth if provided.
ProxyURL string `json:"proxy_url,omitempty"`
// Attributes stores provider specific metadata needed by executors (immutable configuration).
Attributes map[string]string `json:"attributes,omitempty"`
// Metadata stores runtime mutable provider state (e.g. tokens, cookies).
Metadata map[string]any `json:"metadata,omitempty"`
// Quota captures recent quota information for load balancers.
Quota QuotaState `json:"quota"`
// LastError stores the last failure encountered while executing or refreshing.
LastError *Error `json:"last_error,omitempty"`
// CreatedAt is the creation timestamp in UTC.
CreatedAt time.Time `json:"created_at"`
// UpdatedAt is the last modification timestamp in UTC.
UpdatedAt time.Time `json:"updated_at"`
// LastRefreshedAt records the last successful refresh time in UTC.
LastRefreshedAt time.Time `json:"last_refreshed_at"`
// NextRefreshAfter is the earliest time a refresh should retrigger.
NextRefreshAfter time.Time `json:"next_refresh_after"`
// NextRetryAfter is the earliest time a retry should retrigger.
NextRetryAfter time.Time `json:"next_retry_after"`
// ModelStates tracks per-model runtime availability data.
ModelStates map[string]*ModelState `json:"model_states,omitempty"`
// Runtime carries non-serialisable data used during execution (in-memory only).
Runtime any `json:"-"`
}
// QuotaState contains limiter tracking data for a credential.
type QuotaState struct {
// Exceeded indicates the credential recently hit a quota error.
Exceeded bool `json:"exceeded"`
// Reason provides an optional provider specific human readable description.
Reason string `json:"reason,omitempty"`
// NextRecoverAt is when the credential may become available again.
NextRecoverAt time.Time `json:"next_recover_at"`
}
// ModelState captures the execution state for a specific model under an auth entry.
type ModelState struct {
// Status reflects the lifecycle status for this model.
Status Status `json:"status"`
// StatusMessage provides an optional short description of the status.
StatusMessage string `json:"status_message,omitempty"`
// Unavailable mirrors whether the model is temporarily blocked for retries.
Unavailable bool `json:"unavailable"`
// NextRetryAfter defines the per-model retry time.
NextRetryAfter time.Time `json:"next_retry_after"`
// LastError records the latest error observed for this model.
LastError *Error `json:"last_error,omitempty"`
// Quota retains quota information if this model hit rate limits.
Quota QuotaState `json:"quota"`
// UpdatedAt tracks the last update timestamp for this model state.
UpdatedAt time.Time `json:"updated_at"`
}
// Clone shallow copies the Auth structure, duplicating maps to avoid accidental mutation.
func (a *Auth) Clone() *Auth {
if a == nil {
return nil
}
copyAuth := *a
if len(a.Attributes) > 0 {
copyAuth.Attributes = make(map[string]string, len(a.Attributes))
for key, value := range a.Attributes {
copyAuth.Attributes[key] = value
}
}
if len(a.Metadata) > 0 {
copyAuth.Metadata = make(map[string]any, len(a.Metadata))
for key, value := range a.Metadata {
copyAuth.Metadata[key] = value
}
}
if len(a.ModelStates) > 0 {
copyAuth.ModelStates = make(map[string]*ModelState, len(a.ModelStates))
for key, state := range a.ModelStates {
copyAuth.ModelStates[key] = state.Clone()
}
}
copyAuth.Runtime = a.Runtime
return &copyAuth
}
// Clone duplicates a model state including nested error details.
func (m *ModelState) Clone() *ModelState {
if m == nil {
return nil
}
copyState := *m
if m.LastError != nil {
copyState.LastError = &Error{
Code: m.LastError.Code,
Message: m.LastError.Message,
Retryable: m.LastError.Retryable,
HTTPStatus: m.LastError.HTTPStatus,
}
}
return &copyState
}
func (a *Auth) AccountInfo() (string, string) {
if a == nil {
return "", ""
}
if strings.ToLower(a.Provider) == "gemini-web" {
if a.Metadata != nil {
if v, ok := a.Metadata["secure_1psid"].(string); ok && v != "" {
return "cookie", v
}
if v, ok := a.Metadata["__Secure-1PSID"].(string); ok && v != "" {
return "cookie", v
}
}
if a.Attributes != nil {
if v := a.Attributes["secure_1psid"]; v != "" {
return "cookie", v
}
if v := a.Attributes["api_key"]; v != "" {
return "cookie", v
}
}
}
if a.Metadata != nil {
if v, ok := a.Metadata["email"].(string); ok {
return "oauth", v
}
} else if a.Attributes != nil {
if v := a.Attributes["api_key"]; v != "" {
return "api_key", v
}
}
return "", ""
}
// ExpirationTime attempts to extract the credential expiration timestamp from metadata.
// It inspects common keys such as "expired", "expire", "expires_at", and also
// nested "token" objects to remain compatible with legacy auth file formats.
func (a *Auth) ExpirationTime() (time.Time, bool) {
if a == nil {
return time.Time{}, false
}
if ts, ok := expirationFromMap(a.Metadata); ok {
return ts, true
}
return time.Time{}, false
}
var (
refreshLeadMu sync.RWMutex
refreshLeadFactories = make(map[string]func() *time.Duration)
)
func RegisterRefreshLeadProvider(provider string, factory func() *time.Duration) {
provider = strings.ToLower(strings.TrimSpace(provider))
if provider == "" || factory == nil {
return
}
refreshLeadMu.Lock()
refreshLeadFactories[provider] = factory
refreshLeadMu.Unlock()
}
var expireKeys = [...]string{"expired", "expire", "expires_at", "expiresAt", "expiry", "expires"}
func expirationFromMap(meta map[string]any) (time.Time, bool) {
if meta == nil {
return time.Time{}, false
}
for _, key := range expireKeys {
if v, ok := meta[key]; ok {
if ts, ok1 := parseTimeValue(v); ok1 {
return ts, true
}
}
}
for _, nestedKey := range []string{"token", "Token"} {
if nested, ok := meta[nestedKey]; ok {
switch val := nested.(type) {
case map[string]any:
if ts, ok1 := expirationFromMap(val); ok1 {
return ts, true
}
case map[string]string:
temp := make(map[string]any, len(val))
for k, v := range val {
temp[k] = v
}
if ts, ok1 := expirationFromMap(temp); ok1 {
return ts, true
}
}
}
}
return time.Time{}, false
}
func ProviderRefreshLead(provider string, runtime any) *time.Duration {
provider = strings.ToLower(strings.TrimSpace(provider))
if runtime != nil {
if eval, ok := runtime.(interface{ RefreshLead() *time.Duration }); ok {
if lead := eval.RefreshLead(); lead != nil && *lead > 0 {
return lead
}
}
}
refreshLeadMu.RLock()
factory := refreshLeadFactories[provider]
refreshLeadMu.RUnlock()
if factory == nil {
return nil
}
if lead := factory(); lead != nil && *lead > 0 {
return lead
}
return nil
}
func parseTimeValue(v any) (time.Time, bool) {
switch value := v.(type) {
case string:
s := strings.TrimSpace(value)
if s == "" {
return time.Time{}, false
}
layouts := []string{
time.RFC3339,
time.RFC3339Nano,
"2006-01-02 15:04:05",
"2006-01-02T15:04:05Z07:00",
}
for _, layout := range layouts {
if ts, err := time.Parse(layout, s); err == nil {
return ts, true
}
}
if unix, err := strconv.ParseInt(s, 10, 64); err == nil {
return normaliseUnix(unix), true
}
case float64:
return normaliseUnix(int64(value)), true
case int64:
return normaliseUnix(value), true
case json.Number:
if i, err := value.Int64(); err == nil {
return normaliseUnix(i), true
}
if f, err := value.Float64(); err == nil {
return normaliseUnix(int64(f)), true
}
}
return time.Time{}, false
}
func normaliseUnix(raw int64) time.Time {
if raw <= 0 {
return time.Time{}
}
// Heuristic: treat values with millisecond precision (>1e12) accordingly.
if raw > 1_000_000_000_000 {
return time.UnixMilli(raw)
}
return time.Unix(raw, 0)
}

212
sdk/cliproxy/builder.go Normal file
View File

@@ -0,0 +1,212 @@
// Package cliproxy provides the core service implementation for the CLI Proxy API.
// It includes service lifecycle management, authentication handling, file watching,
// and integration with various AI service providers through a unified interface.
package cliproxy
import (
"fmt"
"github.com/router-for-me/CLIProxyAPI/v6/internal/api"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
)
// Builder constructs a Service instance with customizable providers.
// It provides a fluent interface for configuring all aspects of the service
// including authentication, file watching, HTTP server options, and lifecycle hooks.
type Builder struct {
// cfg holds the application configuration.
cfg *config.Config
// configPath is the path to the configuration file.
configPath string
// tokenProvider handles loading token-based clients.
tokenProvider TokenClientProvider
// apiKeyProvider handles loading API key-based clients.
apiKeyProvider APIKeyClientProvider
// watcherFactory creates file watcher instances.
watcherFactory WatcherFactory
// hooks provides lifecycle callbacks.
hooks Hooks
// authManager handles legacy authentication operations.
authManager *sdkAuth.Manager
// accessManager handles request authentication providers.
accessManager *sdkaccess.Manager
// coreManager handles core authentication and execution.
coreManager *coreauth.Manager
// serverOptions contains additional server configuration options.
serverOptions []api.ServerOption
}
// Hooks allows callers to plug into service lifecycle stages.
// These callbacks provide opportunities to perform custom initialization
// and cleanup operations during service startup and shutdown.
type Hooks struct {
// OnBeforeStart is called before the service starts, allowing configuration
// modifications or additional setup.
OnBeforeStart func(*config.Config)
// OnAfterStart is called after the service has started successfully,
// providing access to the service instance for additional operations.
OnAfterStart func(*Service)
}
// NewBuilder creates a Builder with default dependencies left unset.
// Use the fluent interface methods to configure the service before calling Build().
//
// Returns:
// - *Builder: A new builder instance ready for configuration
func NewBuilder() *Builder {
return &Builder{}
}
// WithConfig sets the configuration instance used by the service.
//
// Parameters:
// - cfg: The application configuration
//
// Returns:
// - *Builder: The builder instance for method chaining
func (b *Builder) WithConfig(cfg *config.Config) *Builder {
b.cfg = cfg
return b
}
// WithConfigPath sets the absolute configuration file path used for reload watching.
//
// Parameters:
// - path: The absolute path to the configuration file
//
// Returns:
// - *Builder: The builder instance for method chaining
func (b *Builder) WithConfigPath(path string) *Builder {
b.configPath = path
return b
}
// WithTokenClientProvider overrides the provider responsible for token-backed clients.
func (b *Builder) WithTokenClientProvider(provider TokenClientProvider) *Builder {
b.tokenProvider = provider
return b
}
// WithAPIKeyClientProvider overrides the provider responsible for API key-backed clients.
func (b *Builder) WithAPIKeyClientProvider(provider APIKeyClientProvider) *Builder {
b.apiKeyProvider = provider
return b
}
// WithWatcherFactory allows customizing the watcher factory that handles reloads.
func (b *Builder) WithWatcherFactory(factory WatcherFactory) *Builder {
b.watcherFactory = factory
return b
}
// WithHooks registers lifecycle hooks executed around service startup.
func (b *Builder) WithHooks(h Hooks) *Builder {
b.hooks = h
return b
}
// WithAuthManager overrides the authentication manager used for token lifecycle operations.
func (b *Builder) WithAuthManager(mgr *sdkAuth.Manager) *Builder {
b.authManager = mgr
return b
}
// WithRequestAccessManager overrides the request authentication manager.
func (b *Builder) WithRequestAccessManager(mgr *sdkaccess.Manager) *Builder {
b.accessManager = mgr
return b
}
// WithCoreAuthManager overrides the runtime auth manager responsible for request execution.
func (b *Builder) WithCoreAuthManager(mgr *coreauth.Manager) *Builder {
b.coreManager = mgr
return b
}
// WithServerOptions appends server configuration options used during construction.
func (b *Builder) WithServerOptions(opts ...api.ServerOption) *Builder {
b.serverOptions = append(b.serverOptions, opts...)
return b
}
// Build validates inputs, applies defaults, and returns a ready-to-run service.
func (b *Builder) Build() (*Service, error) {
if b.cfg == nil {
return nil, fmt.Errorf("cliproxy: configuration is required")
}
if b.configPath == "" {
return nil, fmt.Errorf("cliproxy: configuration path is required")
}
tokenProvider := b.tokenProvider
if tokenProvider == nil {
tokenProvider = NewFileTokenClientProvider()
}
apiKeyProvider := b.apiKeyProvider
if apiKeyProvider == nil {
apiKeyProvider = NewAPIKeyClientProvider()
}
watcherFactory := b.watcherFactory
if watcherFactory == nil {
watcherFactory = defaultWatcherFactory
}
authManager := b.authManager
if authManager == nil {
authManager = newDefaultAuthManager()
}
accessManager := b.accessManager
if accessManager == nil {
accessManager = sdkaccess.NewManager()
}
providers, err := sdkaccess.BuildProviders(b.cfg)
if err != nil {
return nil, err
}
accessManager.SetProviders(providers)
coreManager := b.coreManager
if coreManager == nil {
tokenStore := sdkAuth.GetTokenStore()
if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok && b.cfg != nil {
dirSetter.SetBaseDir(b.cfg.AuthDir)
}
store, ok := tokenStore.(coreauth.Store)
if !ok {
return nil, fmt.Errorf("cliproxy: token store does not implement coreauth.Store")
}
coreManager = coreauth.NewManager(store, nil, nil)
}
// Attach a default RoundTripper provider so providers can opt-in per-auth transports.
coreManager.SetRoundTripperProvider(newDefaultRoundTripperProvider())
service := &Service{
cfg: b.cfg,
configPath: b.configPath,
tokenProvider: tokenProvider,
apiKeyProvider: apiKeyProvider,
watcherFactory: watcherFactory,
hooks: b.hooks,
authManager: authManager,
accessManager: accessManager,
coreManager: coreManager,
serverOptions: append([]api.ServerOption(nil), b.serverOptions...),
}
return service, nil
}

View File

@@ -0,0 +1,60 @@
package executor
import (
"net/http"
"net/url"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
)
// Request encapsulates the translated payload that will be sent to a provider executor.
type Request struct {
// Model is the upstream model identifier after translation.
Model string
// Payload is the provider specific JSON payload.
Payload []byte
// Format represents the provider payload schema.
Format sdktranslator.Format
// Metadata carries optional provider specific execution hints.
Metadata map[string]any
}
// Options controls execution behavior for both streaming and non-streaming calls.
type Options struct {
// Stream toggles streaming mode.
Stream bool
// Alt carries optional alternate format hint (e.g. SSE JSON key).
Alt string
// Headers are forwarded to the provider request builder.
Headers http.Header
// Query contains optional query string parameters.
Query url.Values
// OriginalRequest preserves the inbound request bytes prior to translation.
OriginalRequest []byte
// SourceFormat identifies the inbound schema.
SourceFormat sdktranslator.Format
}
// Response wraps either a full provider response or metadata for streaming flows.
type Response struct {
// Payload is the provider response in the executor format.
Payload []byte
// Metadata exposes optional structured data for translators.
Metadata map[string]any
}
// StreamChunk represents a single streaming payload unit emitted by provider executors.
type StreamChunk struct {
// Payload is the raw provider chunk payload.
Payload []byte
// Err reports any terminal error encountered while producing chunks.
Err error
}
// StatusError represents an error that carries an HTTP-like status code.
// Provider executors should implement this when possible to enable
// better auth state updates on failures (e.g., 401/402/429).
type StatusError interface {
error
StatusCode() int
}

View File

@@ -0,0 +1,20 @@
package cliproxy
import "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
// ModelInfo re-exports the registry model info structure.
type ModelInfo = registry.ModelInfo
// ModelRegistry describes registry operations consumed by external callers.
type ModelRegistry interface {
RegisterClient(clientID, clientProvider string, models []*ModelInfo)
UnregisterClient(clientID string)
SetModelQuotaExceeded(clientID, modelID string)
ClearModelQuotaExceeded(clientID, modelID string)
GetAvailableModels(handlerType string) []map[string]any
}
// GlobalModelRegistry returns the shared registry instance.
func GlobalModelRegistry() ModelRegistry {
return registry.GetGlobalRegistry()
}

View File

@@ -0,0 +1,64 @@
package pipeline
import (
"context"
"net/http"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
)
// Context encapsulates execution state shared across middleware, translators, and executors.
type Context struct {
// Request encapsulates the provider facing request payload.
Request cliproxyexecutor.Request
// Options carries execution flags (streaming, headers, etc.).
Options cliproxyexecutor.Options
// Auth references the credential selected for execution.
Auth *cliproxyauth.Auth
// Translator represents the pipeline responsible for schema adaptation.
Translator *sdktranslator.Pipeline
// HTTPClient allows middleware to customise the outbound transport per request.
HTTPClient *http.Client
}
// Hook captures middleware callbacks around execution.
type Hook interface {
BeforeExecute(ctx context.Context, execCtx *Context)
AfterExecute(ctx context.Context, execCtx *Context, resp cliproxyexecutor.Response, err error)
OnStreamChunk(ctx context.Context, execCtx *Context, chunk cliproxyexecutor.StreamChunk)
}
// HookFunc aggregates optional hook implementations.
type HookFunc struct {
Before func(context.Context, *Context)
After func(context.Context, *Context, cliproxyexecutor.Response, error)
Stream func(context.Context, *Context, cliproxyexecutor.StreamChunk)
}
// BeforeExecute implements Hook.
func (h HookFunc) BeforeExecute(ctx context.Context, execCtx *Context) {
if h.Before != nil {
h.Before(ctx, execCtx)
}
}
// AfterExecute implements Hook.
func (h HookFunc) AfterExecute(ctx context.Context, execCtx *Context, resp cliproxyexecutor.Response, err error) {
if h.After != nil {
h.After(ctx, execCtx, resp, err)
}
}
// OnStreamChunk implements Hook.
func (h HookFunc) OnStreamChunk(ctx context.Context, execCtx *Context, chunk cliproxyexecutor.StreamChunk) {
if h.Stream != nil {
h.Stream(ctx, execCtx, chunk)
}
}
// RoundTripperProvider allows injection of custom HTTP transports per auth entry.
type RoundTripperProvider interface {
RoundTripperFor(auth *cliproxyauth.Auth) http.RoundTripper
}

46
sdk/cliproxy/providers.go Normal file
View File

@@ -0,0 +1,46 @@
package cliproxy
import (
"context"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
)
// NewFileTokenClientProvider returns the default token-backed client loader.
func NewFileTokenClientProvider() TokenClientProvider {
return &fileTokenClientProvider{}
}
type fileTokenClientProvider struct{}
func (p *fileTokenClientProvider) Load(ctx context.Context, cfg *config.Config) (*TokenClientResult, error) {
// Stateless executors handle tokens
_ = ctx
_ = cfg
return &TokenClientResult{SuccessfulAuthed: 0}, nil
}
// NewAPIKeyClientProvider returns the default API key client loader that reuses existing logic.
func NewAPIKeyClientProvider() APIKeyClientProvider {
return &apiKeyClientProvider{}
}
type apiKeyClientProvider struct{}
func (p *apiKeyClientProvider) Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error) {
glCount, claudeCount, codexCount, openAICompat := watcher.BuildAPIKeyClients(cfg)
if ctx != nil {
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
}
}
return &APIKeyClientResult{
GeminiKeyCount: glCount,
ClaudeKeyCount: claudeCount,
CodexKeyCount: codexCount,
OpenAICompatCount: openAICompat,
}, nil
}

View File

@@ -0,0 +1,51 @@
package cliproxy
import (
"net/http"
"net/url"
"strings"
"sync"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
)
// defaultRoundTripperProvider returns a per-auth HTTP RoundTripper based on
// the Auth.ProxyURL value. It caches transports per proxy URL string.
type defaultRoundTripperProvider struct {
mu sync.RWMutex
cache map[string]http.RoundTripper
}
func newDefaultRoundTripperProvider() *defaultRoundTripperProvider {
return &defaultRoundTripperProvider{cache: make(map[string]http.RoundTripper)}
}
// RoundTripperFor implements coreauth.RoundTripperProvider.
func (p *defaultRoundTripperProvider) RoundTripperFor(auth *coreauth.Auth) http.RoundTripper {
if auth == nil {
return nil
}
proxy := strings.TrimSpace(auth.ProxyURL)
if proxy == "" {
return nil
}
p.mu.RLock()
rt := p.cache[proxy]
p.mu.RUnlock()
if rt != nil {
return rt
}
// Build HTTP/HTTPS proxy transport; ignore SOCKS for simplicity here.
u, err := url.Parse(proxy)
if err != nil {
return nil
}
if u.Scheme != "http" && u.Scheme != "https" {
return nil
}
transport := &http.Transport{Proxy: http.ProxyURL(u)}
p.mu.Lock()
p.cache[proxy] = transport
p.mu.Unlock()
return transport
}

560
sdk/cliproxy/service.go Normal file
View File

@@ -0,0 +1,560 @@
// Package cliproxy provides the core service implementation for the CLI Proxy API.
// It includes service lifecycle management, authentication handling, file watching,
// and integration with various AI service providers through a unified interface.
package cliproxy
import (
"context"
"errors"
"fmt"
"os"
"strings"
"sync"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/api"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
geminiwebclient "github.com/router-for-me/CLIProxyAPI/v6/internal/provider/gemini-web"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor"
_ "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
"github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
_ "github.com/router-for-me/CLIProxyAPI/v6/sdk/access/providers/configapikey"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
log "github.com/sirupsen/logrus"
)
// Service wraps the proxy server lifecycle so external programs can embed the CLI proxy.
// It manages the complete lifecycle including authentication, file watching, HTTP server,
// and integration with various AI service providers.
type Service struct {
// cfg holds the current application configuration.
cfg *config.Config
// cfgMu protects concurrent access to the configuration.
cfgMu sync.RWMutex
// configPath is the path to the configuration file.
configPath string
// tokenProvider handles loading token-based clients.
tokenProvider TokenClientProvider
// apiKeyProvider handles loading API key-based clients.
apiKeyProvider APIKeyClientProvider
// watcherFactory creates file watcher instances.
watcherFactory WatcherFactory
// hooks provides lifecycle callbacks.
hooks Hooks
// serverOptions contains additional server configuration options.
serverOptions []api.ServerOption
// server is the HTTP API server instance.
server *api.Server
// serverErr channel for server startup/shutdown errors.
serverErr chan error
// watcher handles file system monitoring.
watcher *WatcherWrapper
// watcherCancel cancels the watcher context.
watcherCancel context.CancelFunc
// authUpdates channel for authentication updates.
authUpdates chan watcher.AuthUpdate
// authQueueStop cancels the auth update queue processing.
authQueueStop context.CancelFunc
// authManager handles legacy authentication operations.
authManager *sdkAuth.Manager
// accessManager handles request authentication providers.
accessManager *sdkaccess.Manager
// coreManager handles core authentication and execution.
coreManager *coreauth.Manager
// shutdownOnce ensures shutdown is called only once.
shutdownOnce sync.Once
}
// RegisterUsagePlugin registers a usage plugin on the global usage manager.
// This allows external code to monitor API usage and token consumption.
//
// Parameters:
// - plugin: The usage plugin to register
func (s *Service) RegisterUsagePlugin(plugin usage.Plugin) {
usage.RegisterPlugin(plugin)
}
// newDefaultAuthManager creates a default authentication manager with all supported providers.
func newDefaultAuthManager() *sdkAuth.Manager {
return sdkAuth.NewManager(
sdkAuth.GetTokenStore(),
sdkAuth.NewGeminiAuthenticator(),
sdkAuth.NewCodexAuthenticator(),
sdkAuth.NewClaudeAuthenticator(),
sdkAuth.NewQwenAuthenticator(),
)
}
func (s *Service) refreshAccessProviders(cfg *config.Config) {
if s == nil || s.accessManager == nil || cfg == nil {
return
}
providers, err := sdkaccess.BuildProviders(cfg)
if err != nil {
log.Errorf("failed to rebuild request auth providers: %v", err)
return
}
s.accessManager.SetProviders(providers)
}
func (s *Service) ensureAuthUpdateQueue(ctx context.Context) {
if s == nil {
return
}
if s.authUpdates == nil {
s.authUpdates = make(chan watcher.AuthUpdate, 256)
}
if s.authQueueStop != nil {
return
}
queueCtx, cancel := context.WithCancel(ctx)
s.authQueueStop = cancel
go s.consumeAuthUpdates(queueCtx)
}
func (s *Service) consumeAuthUpdates(ctx context.Context) {
for {
select {
case <-ctx.Done():
return
case update, ok := <-s.authUpdates:
if !ok {
return
}
s.handleAuthUpdate(ctx, update)
labelDrain:
for {
select {
case nextUpdate := <-s.authUpdates:
s.handleAuthUpdate(ctx, nextUpdate)
default:
break labelDrain
}
}
}
}
}
func (s *Service) handleAuthUpdate(ctx context.Context, update watcher.AuthUpdate) {
if s == nil {
return
}
s.cfgMu.RLock()
cfg := s.cfg
s.cfgMu.RUnlock()
if cfg == nil || s.coreManager == nil {
return
}
switch update.Action {
case watcher.AuthUpdateActionAdd, watcher.AuthUpdateActionModify:
if update.Auth == nil || update.Auth.ID == "" {
return
}
s.applyCoreAuthAddOrUpdate(ctx, update.Auth)
case watcher.AuthUpdateActionDelete:
id := update.ID
if id == "" && update.Auth != nil {
id = update.Auth.ID
}
if id == "" {
return
}
s.applyCoreAuthRemoval(ctx, id)
default:
log.Debugf("received unknown auth update action: %v", update.Action)
}
}
func (s *Service) applyCoreAuthAddOrUpdate(ctx context.Context, auth *coreauth.Auth) {
if s == nil || auth == nil || auth.ID == "" {
return
}
if s.coreManager == nil {
return
}
auth = auth.Clone()
s.ensureExecutorsForAuth(auth)
s.registerModelsForAuth(auth)
if existing, ok := s.coreManager.GetByID(auth.ID); ok && existing != nil {
auth.CreatedAt = existing.CreatedAt
auth.LastRefreshedAt = existing.LastRefreshedAt
auth.NextRefreshAfter = existing.NextRefreshAfter
if _, err := s.coreManager.Update(ctx, auth); err != nil {
log.Errorf("failed to update auth %s: %v", auth.ID, err)
}
return
}
if _, err := s.coreManager.Register(ctx, auth); err != nil {
log.Errorf("failed to register auth %s: %v", auth.ID, err)
}
}
func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) {
if s == nil || id == "" {
return
}
if s.coreManager == nil {
return
}
GlobalModelRegistry().UnregisterClient(id)
if existing, ok := s.coreManager.GetByID(id); ok && existing != nil {
existing.Disabled = true
existing.Status = coreauth.StatusDisabled
if _, err := s.coreManager.Update(ctx, existing); err != nil {
log.Errorf("failed to disable auth %s: %v", id, err)
}
}
}
func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) {
if s == nil || a == nil {
return
}
switch strings.ToLower(a.Provider) {
case "gemini":
s.coreManager.RegisterExecutor(executor.NewGeminiExecutor(s.cfg))
case "gemini-cli":
s.coreManager.RegisterExecutor(executor.NewGeminiCLIExecutor(s.cfg))
case "gemini-web":
s.coreManager.RegisterExecutor(executor.NewGeminiWebExecutor(s.cfg))
case "claude":
s.coreManager.RegisterExecutor(executor.NewClaudeExecutor(s.cfg))
case "codex":
s.coreManager.RegisterExecutor(executor.NewCodexExecutor(s.cfg))
case "qwen":
s.coreManager.RegisterExecutor(executor.NewQwenExecutor(s.cfg))
default:
providerKey := strings.ToLower(strings.TrimSpace(a.Provider))
if providerKey == "" {
providerKey = "openai-compatibility"
}
s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(providerKey, s.cfg))
}
}
// Run starts the service and blocks until the context is cancelled or the server stops.
// It initializes all components including authentication, file watching, HTTP server,
// and starts processing requests. The method blocks until the context is cancelled.
//
// Parameters:
// - ctx: The context for controlling the service lifecycle
//
// Returns:
// - error: An error if the service fails to start or run
func (s *Service) Run(ctx context.Context) error {
if s == nil {
return fmt.Errorf("cliproxy: service is nil")
}
if ctx == nil {
ctx = context.Background()
}
usage.StartDefault(ctx)
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
defer shutdownCancel()
defer func() {
if err := s.Shutdown(shutdownCtx); err != nil {
log.Errorf("service shutdown returned error: %v", err)
}
}()
if err := s.ensureAuthDir(); err != nil {
return err
}
if s.coreManager != nil {
if errLoad := s.coreManager.Load(ctx); errLoad != nil {
log.Warnf("failed to load auth store: %v", errLoad)
}
}
tokenResult, err := s.tokenProvider.Load(ctx, s.cfg)
if err != nil && !errors.Is(err, context.Canceled) {
return err
}
if tokenResult == nil {
tokenResult = &TokenClientResult{}
}
apiKeyResult, err := s.apiKeyProvider.Load(ctx, s.cfg)
if err != nil && !errors.Is(err, context.Canceled) {
return err
}
if apiKeyResult == nil {
apiKeyResult = &APIKeyClientResult{}
}
// legacy clients removed; no caches to refresh
// handlers no longer depend on legacy clients; pass nil slice initially
s.refreshAccessProviders(s.cfg)
s.server = api.NewServer(s.cfg, s.coreManager, s.accessManager, s.configPath, s.serverOptions...)
if s.authManager == nil {
s.authManager = newDefaultAuthManager()
}
if s.hooks.OnBeforeStart != nil {
s.hooks.OnBeforeStart(s.cfg)
}
s.serverErr = make(chan error, 1)
go func() {
if errStart := s.server.Start(); errStart != nil {
s.serverErr <- errStart
} else {
s.serverErr <- nil
}
}()
time.Sleep(100 * time.Millisecond)
log.Info("API server started successfully")
if s.hooks.OnAfterStart != nil {
s.hooks.OnAfterStart(s)
}
var watcherWrapper *WatcherWrapper
reloadCallback := func(newCfg *config.Config) {
if newCfg == nil {
s.cfgMu.RLock()
newCfg = s.cfg
s.cfgMu.RUnlock()
}
if newCfg == nil {
return
}
s.refreshAccessProviders(newCfg)
if s.server != nil {
s.server.UpdateClients(newCfg)
}
s.cfgMu.Lock()
s.cfg = newCfg
s.cfgMu.Unlock()
}
watcherWrapper, err = s.watcherFactory(s.configPath, s.cfg.AuthDir, reloadCallback)
if err != nil {
return fmt.Errorf("cliproxy: failed to create watcher: %w", err)
}
s.watcher = watcherWrapper
s.ensureAuthUpdateQueue(ctx)
if s.authUpdates != nil {
watcherWrapper.SetAuthUpdateQueue(s.authUpdates)
}
watcherWrapper.SetConfig(s.cfg)
watcherCtx, watcherCancel := context.WithCancel(context.Background())
s.watcherCancel = watcherCancel
if err = watcherWrapper.Start(watcherCtx); err != nil {
return fmt.Errorf("cliproxy: failed to start watcher: %w", err)
}
log.Info("file watcher started for config and auth directory changes")
// Prefer core auth manager auto refresh if available.
if s.coreManager != nil {
interval := 15 * time.Minute
s.coreManager.StartAutoRefresh(context.Background(), interval)
log.Infof("core auth auto-refresh started (interval=%s)", interval)
}
authFileCount := util.CountAuthFiles(s.cfg.AuthDir)
totalNewClients := authFileCount + apiKeyResult.GeminiKeyCount + apiKeyResult.ClaudeKeyCount + apiKeyResult.CodexKeyCount + apiKeyResult.OpenAICompatCount
log.Infof("full client load complete - %d clients (%d auth files + %d GL API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)",
totalNewClients,
authFileCount,
apiKeyResult.GeminiKeyCount,
apiKeyResult.ClaudeKeyCount,
apiKeyResult.CodexKeyCount,
apiKeyResult.OpenAICompatCount,
)
select {
case <-ctx.Done():
log.Debug("service context cancelled, shutting down...")
return ctx.Err()
case err = <-s.serverErr:
return err
}
}
// Shutdown gracefully stops background workers and the HTTP server.
// It ensures all resources are properly cleaned up and connections are closed.
// The shutdown is idempotent and can be called multiple times safely.
//
// Parameters:
// - ctx: The context for controlling the shutdown timeout
//
// Returns:
// - error: An error if shutdown fails
func (s *Service) Shutdown(ctx context.Context) error {
if s == nil {
return nil
}
var shutdownErr error
s.shutdownOnce.Do(func() {
if ctx == nil {
ctx = context.Background()
}
// legacy refresh loop removed; only stopping core auth manager below
if s.watcherCancel != nil {
s.watcherCancel()
}
if s.coreManager != nil {
s.coreManager.StopAutoRefresh()
}
if s.watcher != nil {
if err := s.watcher.Stop(); err != nil {
log.Errorf("failed to stop file watcher: %v", err)
shutdownErr = err
}
}
if s.authQueueStop != nil {
s.authQueueStop()
s.authQueueStop = nil
}
// no legacy clients to persist
if s.server != nil {
shutdownCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
if err := s.server.Stop(shutdownCtx); err != nil {
log.Errorf("error stopping API server: %v", err)
if shutdownErr == nil {
shutdownErr = err
}
}
}
usage.StopDefault()
})
return shutdownErr
}
func (s *Service) ensureAuthDir() error {
info, err := os.Stat(s.cfg.AuthDir)
if err != nil {
if os.IsNotExist(err) {
if mkErr := os.MkdirAll(s.cfg.AuthDir, 0o755); mkErr != nil {
return fmt.Errorf("cliproxy: failed to create auth directory %s: %w", s.cfg.AuthDir, mkErr)
}
log.Infof("created missing auth directory: %s", s.cfg.AuthDir)
return nil
}
return fmt.Errorf("cliproxy: error checking auth directory %s: %w", s.cfg.AuthDir, err)
}
if !info.IsDir() {
return fmt.Errorf("cliproxy: auth path exists but is not a directory: %s", s.cfg.AuthDir)
}
return nil
}
// registerModelsForAuth (re)binds provider models in the global registry using the core auth ID as client identifier.
func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
if a == nil || a.ID == "" {
return
}
// Unregister legacy client ID (if present) to avoid double counting
if a.Runtime != nil {
if idGetter, ok := a.Runtime.(interface{ GetClientID() string }); ok {
if rid := idGetter.GetClientID(); rid != "" && rid != a.ID {
GlobalModelRegistry().UnregisterClient(rid)
}
}
}
provider := strings.ToLower(strings.TrimSpace(a.Provider))
var models []*ModelInfo
switch provider {
case "gemini":
models = registry.GetGeminiModels()
case "gemini-cli":
models = registry.GetGeminiCLIModels()
case "gemini-web":
models = geminiwebclient.GetGeminiWebAliasedModels()
case "claude":
models = registry.GetClaudeModels()
case "codex":
models = registry.GetOpenAIModels()
case "qwen":
models = registry.GetQwenModels()
default:
// Handle OpenAI-compatibility providers by name using config
if s.cfg != nil {
providerKey := provider
compatName := strings.TrimSpace(a.Provider)
if strings.EqualFold(providerKey, "openai-compatibility") {
if a.Attributes != nil {
if v := strings.TrimSpace(a.Attributes["compat_name"]); v != "" {
compatName = v
}
if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" {
providerKey = strings.ToLower(v)
}
}
if providerKey == "openai-compatibility" && compatName != "" {
providerKey = strings.ToLower(compatName)
}
}
for i := range s.cfg.OpenAICompatibility {
compat := &s.cfg.OpenAICompatibility[i]
if strings.EqualFold(compat.Name, compatName) {
// Convert compatibility models to registry models
ms := make([]*ModelInfo, 0, len(compat.Models))
for j := range compat.Models {
m := compat.Models[j]
ms = append(ms, &ModelInfo{
ID: m.Alias,
Object: "model",
Created: time.Now().Unix(),
OwnedBy: compat.Name,
Type: "openai-compatibility",
DisplayName: m.Name,
})
}
// Register and return
if len(ms) > 0 {
if providerKey == "" {
providerKey = "openai-compatibility"
}
GlobalModelRegistry().RegisterClient(a.ID, providerKey, ms)
}
return
}
}
}
}
if len(models) > 0 {
key := provider
if key == "" {
key = strings.ToLower(strings.TrimSpace(a.Provider))
}
GlobalModelRegistry().RegisterClient(a.ID, key, models)
}
}

135
sdk/cliproxy/types.go Normal file
View File

@@ -0,0 +1,135 @@
// Package cliproxy provides the core service implementation for the CLI Proxy API.
// It includes service lifecycle management, authentication handling, file watching,
// and integration with various AI service providers through a unified interface.
package cliproxy
import (
"context"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
)
// TokenClientProvider loads clients backed by stored authentication tokens.
// It provides an interface for loading authentication tokens from various sources
// and creating clients for AI service providers.
type TokenClientProvider interface {
// Load loads token-based clients from the configured source.
//
// Parameters:
// - ctx: The context for the loading operation
// - cfg: The application configuration
//
// Returns:
// - *TokenClientResult: The result containing loaded clients
// - error: An error if loading fails
Load(ctx context.Context, cfg *config.Config) (*TokenClientResult, error)
}
// TokenClientResult represents clients generated from persisted tokens.
// It contains metadata about the loading operation and the number of successful authentications.
type TokenClientResult struct {
// SuccessfulAuthed is the number of successfully authenticated clients.
SuccessfulAuthed int
}
// APIKeyClientProvider loads clients backed directly by configured API keys.
// It provides an interface for loading API key-based clients for various AI service providers.
type APIKeyClientProvider interface {
// Load loads API key-based clients from the configuration.
//
// Parameters:
// - ctx: The context for the loading operation
// - cfg: The application configuration
//
// Returns:
// - *APIKeyClientResult: The result containing loaded clients
// - error: An error if loading fails
Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error)
}
// APIKeyClientResult contains API key based clients along with type counts.
// It provides metadata about the number of clients loaded for each provider type.
type APIKeyClientResult struct {
// GeminiKeyCount is the number of Gemini API key clients loaded.
GeminiKeyCount int
// ClaudeKeyCount is the number of Claude API key clients loaded.
ClaudeKeyCount int
// CodexKeyCount is the number of Codex API key clients loaded.
CodexKeyCount int
// OpenAICompatCount is the number of OpenAI-compatible API key clients loaded.
OpenAICompatCount int
}
// WatcherFactory creates a watcher for configuration and token changes.
// The reload callback receives the updated configuration when changes are detected.
//
// Parameters:
// - configPath: The path to the configuration file to watch
// - authDir: The directory containing authentication tokens to watch
// - reload: The callback function to call when changes are detected
//
// Returns:
// - *WatcherWrapper: A watcher wrapper instance
// - error: An error if watcher creation fails
type WatcherFactory func(configPath, authDir string, reload func(*config.Config)) (*WatcherWrapper, error)
// WatcherWrapper exposes the subset of watcher methods required by the SDK.
type WatcherWrapper struct {
start func(ctx context.Context) error
stop func() error
setConfig func(cfg *config.Config)
snapshotAuths func() []*coreauth.Auth
setUpdateQueue func(queue chan<- watcher.AuthUpdate)
}
// Start proxies to the underlying watcher Start implementation.
func (w *WatcherWrapper) Start(ctx context.Context) error {
if w == nil || w.start == nil {
return nil
}
return w.start(ctx)
}
// Stop proxies to the underlying watcher Stop implementation.
func (w *WatcherWrapper) Stop() error {
if w == nil || w.stop == nil {
return nil
}
return w.stop()
}
// SetConfig updates the watcher configuration cache.
func (w *WatcherWrapper) SetConfig(cfg *config.Config) {
if w == nil || w.setConfig == nil {
return
}
w.setConfig(cfg)
}
// SetClients updates the watcher file-backed clients registry.
// SetClients and SetAPIKeyClients removed; watcher manages its own caches
// SnapshotClients returns the current combined clients snapshot from the underlying watcher.
// SnapshotClients removed; use SnapshotAuths
// SnapshotAuths returns the current auth entries derived from legacy clients.
func (w *WatcherWrapper) SnapshotAuths() []*coreauth.Auth {
if w == nil || w.snapshotAuths == nil {
return nil
}
return w.snapshotAuths()
}
// SetAuthUpdateQueue registers the channel used to propagate auth updates.
func (w *WatcherWrapper) SetAuthUpdateQueue(queue chan<- watcher.AuthUpdate) {
if w == nil || w.setUpdateQueue == nil {
return
}
w.setUpdateQueue(queue)
}

View File

@@ -0,0 +1,178 @@
package usage
import (
"context"
"sync"
"time"
log "github.com/sirupsen/logrus"
)
// Record contains the usage statistics captured for a single provider request.
type Record struct {
Provider string
Model string
APIKey string
AuthID string
RequestedAt time.Time
Detail Detail
}
// Detail holds the token usage breakdown.
type Detail struct {
InputTokens int64
OutputTokens int64
ReasoningTokens int64
CachedTokens int64
TotalTokens int64
}
// Plugin consumes usage records emitted by the proxy runtime.
type Plugin interface {
HandleUsage(ctx context.Context, record Record)
}
type queueItem struct {
ctx context.Context
record Record
}
// Manager maintains a queue of usage records and delivers them to registered plugins.
type Manager struct {
once sync.Once
stopOnce sync.Once
cancel context.CancelFunc
mu sync.Mutex
cond *sync.Cond
queue []queueItem
closed bool
pluginsMu sync.RWMutex
plugins []Plugin
}
// NewManager constructs a manager with a buffered queue.
func NewManager(buffer int) *Manager {
m := &Manager{}
m.cond = sync.NewCond(&m.mu)
return m
}
// Start launches the background dispatcher. Calling Start multiple times is safe.
func (m *Manager) Start(ctx context.Context) {
if m == nil {
return
}
m.once.Do(func() {
if ctx == nil {
ctx = context.Background()
}
var workerCtx context.Context
workerCtx, m.cancel = context.WithCancel(ctx)
go m.run(workerCtx)
})
}
// Stop stops the dispatcher and drains the queue.
func (m *Manager) Stop() {
if m == nil {
return
}
m.stopOnce.Do(func() {
if m.cancel != nil {
m.cancel()
}
m.mu.Lock()
m.closed = true
m.mu.Unlock()
m.cond.Broadcast()
})
}
// Register appends a plugin to the delivery list.
func (m *Manager) Register(plugin Plugin) {
if m == nil || plugin == nil {
return
}
m.pluginsMu.Lock()
m.plugins = append(m.plugins, plugin)
m.pluginsMu.Unlock()
}
// Publish enqueues a usage record for processing. If no plugin is registered
// the record will be discarded downstream.
func (m *Manager) Publish(ctx context.Context, record Record) {
if m == nil {
return
}
// ensure worker is running even if Start was not called explicitly
m.Start(context.Background())
m.mu.Lock()
if m.closed {
m.mu.Unlock()
return
}
m.queue = append(m.queue, queueItem{ctx: ctx, record: record})
m.mu.Unlock()
m.cond.Signal()
}
func (m *Manager) run(ctx context.Context) {
for {
m.mu.Lock()
for !m.closed && len(m.queue) == 0 {
m.cond.Wait()
}
if len(m.queue) == 0 && m.closed {
m.mu.Unlock()
return
}
item := m.queue[0]
m.queue = m.queue[1:]
m.mu.Unlock()
m.dispatch(item)
}
}
func (m *Manager) dispatch(item queueItem) {
m.pluginsMu.RLock()
plugins := make([]Plugin, len(m.plugins))
copy(plugins, m.plugins)
m.pluginsMu.RUnlock()
if len(plugins) == 0 {
return
}
for _, plugin := range plugins {
if plugin == nil {
continue
}
safeInvoke(plugin, item.ctx, item.record)
}
}
func safeInvoke(plugin Plugin, ctx context.Context, record Record) {
defer func() {
if r := recover(); r != nil {
log.Errorf("usage: plugin panic recovered: %v", r)
}
}()
plugin.HandleUsage(ctx, record)
}
var defaultManager = NewManager(512)
// DefaultManager returns the global usage manager instance.
func DefaultManager() *Manager { return defaultManager }
// RegisterPlugin registers a plugin on the default manager.
func RegisterPlugin(plugin Plugin) { DefaultManager().Register(plugin) }
// PublishRecord publishes a record using the default manager.
func PublishRecord(ctx context.Context, record Record) { DefaultManager().Publish(ctx, record) }
// StartDefault starts the default manager's dispatcher.
func StartDefault(ctx context.Context) { DefaultManager().Start(ctx) }
// StopDefault stops the default manager's dispatcher.
func StopDefault() { DefaultManager().Stop() }

32
sdk/cliproxy/watcher.go Normal file
View File

@@ -0,0 +1,32 @@
package cliproxy
import (
"context"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/watcher"
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
)
func defaultWatcherFactory(configPath, authDir string, reload func(*config.Config)) (*WatcherWrapper, error) {
w, err := watcher.NewWatcher(configPath, authDir, reload)
if err != nil {
return nil, err
}
return &WatcherWrapper{
start: func(ctx context.Context) error {
return w.Start(ctx)
},
stop: func() error {
return w.Stop()
},
setConfig: func(cfg *config.Config) {
w.SetConfig(cfg)
},
snapshotAuths: func() []*coreauth.Auth { return w.SnapshotCoreAuths() },
setUpdateQueue: func(queue chan<- watcher.AuthUpdate) {
w.SetAuthUpdateQueue(queue)
},
}, nil
}