mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-19 04:40:52 +08:00
refactor(config): consolidate Amp settings into AmpCode struct
This commit is contained in:
@@ -95,7 +95,8 @@ func (m *AmpModule) Name() string {
|
|||||||
// This implements the RouteModuleV2 interface with Context.
|
// This implements the RouteModuleV2 interface with Context.
|
||||||
// Routes are registered only once via sync.Once for idempotent behavior.
|
// Routes are registered only once via sync.Once for idempotent behavior.
|
||||||
func (m *AmpModule) Register(ctx modules.Context) error {
|
func (m *AmpModule) Register(ctx modules.Context) error {
|
||||||
upstreamURL := strings.TrimSpace(ctx.Config.AmpUpstreamURL)
|
settings := ctx.Config.AmpCode
|
||||||
|
upstreamURL := strings.TrimSpace(settings.UpstreamURL)
|
||||||
|
|
||||||
// Determine auth middleware (from module or context)
|
// Determine auth middleware (from module or context)
|
||||||
auth := m.getAuthMiddleware(ctx)
|
auth := m.getAuthMiddleware(ctx)
|
||||||
@@ -104,7 +105,7 @@ func (m *AmpModule) Register(ctx modules.Context) error {
|
|||||||
var regErr error
|
var regErr error
|
||||||
m.registerOnce.Do(func() {
|
m.registerOnce.Do(func() {
|
||||||
// Initialize model mapper from config (for routing unavailable models to alternatives)
|
// Initialize model mapper from config (for routing unavailable models to alternatives)
|
||||||
m.modelMapper = NewModelMapper(ctx.Config.AmpModelMappings)
|
m.modelMapper = NewModelMapper(settings.ModelMappings)
|
||||||
|
|
||||||
// Always register provider aliases - these work without an upstream
|
// Always register provider aliases - these work without an upstream
|
||||||
m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)
|
m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)
|
||||||
@@ -120,7 +121,7 @@ func (m *AmpModule) Register(ctx modules.Context) error {
|
|||||||
// Create secret source with precedence: config > env > file
|
// Create secret source with precedence: config > env > file
|
||||||
// Cache secrets for 5 minutes to reduce file I/O
|
// Cache secrets for 5 minutes to reduce file I/O
|
||||||
if m.secretSource == nil {
|
if m.secretSource == nil {
|
||||||
m.secretSource = NewMultiSourceSecret(ctx.Config.AmpUpstreamAPIKey, 0 /* default 5min */)
|
m.secretSource = NewMultiSourceSecret(settings.UpstreamAPIKey, 0 /* default 5min */)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Create reverse proxy with gzip handling via ModifyResponse
|
// Create reverse proxy with gzip handling via ModifyResponse
|
||||||
@@ -136,7 +137,7 @@ func (m *AmpModule) Register(ctx modules.Context) error {
|
|||||||
// Register management proxy routes (requires upstream)
|
// Register management proxy routes (requires upstream)
|
||||||
// Restrict to localhost by default for security (prevents drive-by browser attacks)
|
// Restrict to localhost by default for security (prevents drive-by browser attacks)
|
||||||
handler := proxyHandler(proxy)
|
handler := proxyHandler(proxy)
|
||||||
m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler, handler, ctx.Config.AmpRestrictManagementToLocalhost)
|
m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler, handler, settings.RestrictManagementToLocalhost)
|
||||||
|
|
||||||
log.Infof("Amp upstream proxy enabled for: %s", upstreamURL)
|
log.Infof("Amp upstream proxy enabled for: %s", upstreamURL)
|
||||||
log.Debug("Amp provider alias routes registered")
|
log.Debug("Amp provider alias routes registered")
|
||||||
@@ -166,8 +167,9 @@ func (m *AmpModule) getAuthMiddleware(ctx modules.Context) gin.HandlerFunc {
|
|||||||
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
||||||
// Update model mappings (hot-reload supported)
|
// Update model mappings (hot-reload supported)
|
||||||
if m.modelMapper != nil {
|
if m.modelMapper != nil {
|
||||||
log.Infof("amp config updated: reloading %d model mapping(s)", len(cfg.AmpModelMappings))
|
settings := cfg.AmpCode
|
||||||
m.modelMapper.UpdateMappings(cfg.AmpModelMappings)
|
log.Infof("amp config updated: reloading %d model mapping(s)", len(settings.ModelMappings))
|
||||||
|
m.modelMapper.UpdateMappings(settings.ModelMappings)
|
||||||
} else {
|
} else {
|
||||||
log.Warnf("amp model mapper not initialized, skipping model mapping update")
|
log.Warnf("amp model mapper not initialized, skipping model mapping update")
|
||||||
}
|
}
|
||||||
@@ -177,7 +179,7 @@ func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
upstreamURL := strings.TrimSpace(cfg.AmpUpstreamURL)
|
upstreamURL := strings.TrimSpace(cfg.AmpCode.UpstreamURL)
|
||||||
if upstreamURL == "" {
|
if upstreamURL == "" {
|
||||||
log.Warn("Amp upstream URL removed from config, restart required to disable")
|
log.Warn("Amp upstream URL removed from config, restart required to disable")
|
||||||
return nil
|
return nil
|
||||||
|
|||||||
@@ -56,8 +56,10 @@ func TestAmpModule_Register_WithUpstream(t *testing.T) {
|
|||||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||||
|
|
||||||
cfg := &config.Config{
|
cfg := &config.Config{
|
||||||
AmpUpstreamURL: upstream.URL,
|
AmpCode: config.AmpCode{
|
||||||
AmpUpstreamAPIKey: "test-key",
|
UpstreamURL: upstream.URL,
|
||||||
|
UpstreamAPIKey: "test-key",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||||
@@ -86,7 +88,9 @@ func TestAmpModule_Register_WithoutUpstream(t *testing.T) {
|
|||||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||||
|
|
||||||
cfg := &config.Config{
|
cfg := &config.Config{
|
||||||
AmpUpstreamURL: "", // No upstream
|
AmpCode: config.AmpCode{
|
||||||
|
UpstreamURL: "", // No upstream
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||||
@@ -121,7 +125,9 @@ func TestAmpModule_Register_InvalidUpstream(t *testing.T) {
|
|||||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||||
|
|
||||||
cfg := &config.Config{
|
cfg := &config.Config{
|
||||||
AmpUpstreamURL: "://invalid-url",
|
AmpCode: config.AmpCode{
|
||||||
|
UpstreamURL: "://invalid-url",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||||
@@ -151,7 +157,7 @@ func TestAmpModule_OnConfigUpdated_CacheInvalidation(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Update config - should invalidate cache
|
// Update config - should invalidate cache
|
||||||
if err := m.OnConfigUpdated(&config.Config{AmpUpstreamURL: "http://x"}); err != nil {
|
if err := m.OnConfigUpdated(&config.Config{AmpCode: config.AmpCode{UpstreamURL: "http://x"}}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -175,7 +181,7 @@ func TestAmpModule_OnConfigUpdated_URLRemoved(t *testing.T) {
|
|||||||
m.secretSource = ms
|
m.secretSource = ms
|
||||||
|
|
||||||
// Config update with empty URL - should log warning but not error
|
// Config update with empty URL - should log warning but not error
|
||||||
cfg := &config.Config{AmpUpstreamURL: ""}
|
cfg := &config.Config{AmpCode: config.AmpCode{UpstreamURL: ""}}
|
||||||
|
|
||||||
if err := m.OnConfigUpdated(cfg); err != nil {
|
if err := m.OnConfigUpdated(cfg); err != nil {
|
||||||
t.Fatalf("unexpected error: %v", err)
|
t.Fatalf("unexpected error: %v", err)
|
||||||
@@ -187,7 +193,7 @@ func TestAmpModule_OnConfigUpdated_NonMultiSourceSecret(t *testing.T) {
|
|||||||
m := &AmpModule{enabled: true}
|
m := &AmpModule{enabled: true}
|
||||||
m.secretSource = NewStaticSecretSource("static-key")
|
m.secretSource = NewStaticSecretSource("static-key")
|
||||||
|
|
||||||
cfg := &config.Config{AmpUpstreamURL: "http://example.com"}
|
cfg := &config.Config{AmpCode: config.AmpCode{UpstreamURL: "http://example.com"}}
|
||||||
|
|
||||||
// Should not error or panic
|
// Should not error or panic
|
||||||
if err := m.OnConfigUpdated(cfg); err != nil {
|
if err := m.OnConfigUpdated(cfg); err != nil {
|
||||||
@@ -240,8 +246,10 @@ func TestAmpModule_SecretSource_FromConfig(t *testing.T) {
|
|||||||
|
|
||||||
// Config with explicit API key
|
// Config with explicit API key
|
||||||
cfg := &config.Config{
|
cfg := &config.Config{
|
||||||
AmpUpstreamURL: upstream.URL,
|
AmpCode: config.AmpCode{
|
||||||
AmpUpstreamAPIKey: "config-key",
|
UpstreamURL: upstream.URL,
|
||||||
|
UpstreamAPIKey: "config-key",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||||
@@ -283,7 +291,7 @@ func TestAmpModule_ProviderAliasesAlwaysRegistered(t *testing.T) {
|
|||||||
|
|
||||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||||
|
|
||||||
cfg := &config.Config{AmpUpstreamURL: scenario.configURL}
|
cfg := &config.Config{AmpCode: config.AmpCode{UpstreamURL: scenario.configURL}}
|
||||||
|
|
||||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||||
if err := m.Register(ctx); err != nil && scenario.configURL != "" {
|
if err := m.Register(ctx); err != nil && scenario.configURL != "" {
|
||||||
|
|||||||
@@ -26,22 +26,8 @@ type Config struct {
|
|||||||
// TLS config controls HTTPS server settings.
|
// TLS config controls HTTPS server settings.
|
||||||
TLS TLSConfig `yaml:"tls" json:"tls"`
|
TLS TLSConfig `yaml:"tls" json:"tls"`
|
||||||
|
|
||||||
// AmpUpstreamURL defines the upstream Amp control plane used for non-provider calls.
|
// RemoteManagement nests management-related options under 'remote-management'.
|
||||||
AmpUpstreamURL string `yaml:"amp-upstream-url" json:"amp-upstream-url"`
|
RemoteManagement RemoteManagement `yaml:"remote-management" json:"-"`
|
||||||
|
|
||||||
// AmpUpstreamAPIKey optionally overrides the Authorization header when proxying Amp upstream calls.
|
|
||||||
AmpUpstreamAPIKey string `yaml:"amp-upstream-api-key" json:"amp-upstream-api-key"`
|
|
||||||
|
|
||||||
// AmpRestrictManagementToLocalhost restricts Amp management routes (/api/user, /api/threads, etc.)
|
|
||||||
// to only accept connections from localhost (127.0.0.1, ::1). When true, prevents drive-by
|
|
||||||
// browser attacks and remote access to management endpoints. Default: true (recommended).
|
|
||||||
AmpRestrictManagementToLocalhost bool `yaml:"amp-restrict-management-to-localhost" json:"amp-restrict-management-to-localhost"`
|
|
||||||
|
|
||||||
// AmpModelMappings defines model name mappings for Amp CLI requests.
|
|
||||||
// When Amp requests a model that isn't available locally, these mappings
|
|
||||||
// allow routing to an alternative model that IS available.
|
|
||||||
// Example: Map "claude-opus-4.5" -> "claude-sonnet-4" when opus isn't available.
|
|
||||||
AmpModelMappings []AmpModelMapping `yaml:"amp-model-mappings" json:"amp-model-mappings"`
|
|
||||||
|
|
||||||
// AuthDir is the directory where authentication token files are stored.
|
// AuthDir is the directory where authentication token files are stored.
|
||||||
AuthDir string `yaml:"auth-dir" json:"-"`
|
AuthDir string `yaml:"auth-dir" json:"-"`
|
||||||
@@ -58,44 +44,44 @@ type Config struct {
|
|||||||
// DisableCooling disables quota cooldown scheduling when true.
|
// DisableCooling disables quota cooldown scheduling when true.
|
||||||
DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
|
DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
|
||||||
|
|
||||||
|
// RequestRetry defines the retry times when the request failed.
|
||||||
|
RequestRetry int `yaml:"request-retry" json:"request-retry"`
|
||||||
|
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
|
||||||
|
MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"`
|
||||||
|
|
||||||
// QuotaExceeded defines the behavior when a quota is exceeded.
|
// QuotaExceeded defines the behavior when a quota is exceeded.
|
||||||
QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
|
QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
|
||||||
|
|
||||||
// WebsocketAuth enables or disables authentication for the WebSocket API.
|
// WebsocketAuth enables or disables authentication for the WebSocket API.
|
||||||
WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
|
WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
|
||||||
|
|
||||||
|
// GeminiKey defines Gemini API key configurations with optional routing overrides.
|
||||||
|
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
|
||||||
|
|
||||||
// GlAPIKey exposes the legacy generative language API key list for backward compatibility.
|
// GlAPIKey exposes the legacy generative language API key list for backward compatibility.
|
||||||
GlAPIKey []string `yaml:"generative-language-api-key" json:"generative-language-api-key"`
|
GlAPIKey []string `yaml:"generative-language-api-key" json:"generative-language-api-key"`
|
||||||
|
|
||||||
// GeminiKey defines Gemini API key configurations with optional routing overrides.
|
// Codex defines a list of Codex API key configurations as specified in the YAML configuration file.
|
||||||
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
|
CodexKey []CodexKey `yaml:"codex-api-key" json:"codex-api-key"`
|
||||||
|
|
||||||
|
// ClaudeKey defines a list of Claude API key configurations as specified in the YAML configuration file.
|
||||||
|
ClaudeKey []ClaudeKey `yaml:"claude-api-key" json:"claude-api-key"`
|
||||||
|
|
||||||
|
// OpenAICompatibility defines OpenAI API compatibility configurations for external providers.
|
||||||
|
OpenAICompatibility []OpenAICompatibility `yaml:"openai-compatibility" json:"openai-compatibility"`
|
||||||
|
|
||||||
// VertexCompatAPIKey defines Vertex AI-compatible API key configurations for third-party providers.
|
// VertexCompatAPIKey defines Vertex AI-compatible API key configurations for third-party providers.
|
||||||
// Used for services that use Vertex AI-style paths but with simple API key authentication.
|
// Used for services that use Vertex AI-style paths but with simple API key authentication.
|
||||||
VertexCompatAPIKey []VertexCompatKey `yaml:"vertex-api-key" json:"vertex-api-key"`
|
VertexCompatAPIKey []VertexCompatKey `yaml:"vertex-api-key" json:"vertex-api-key"`
|
||||||
|
|
||||||
// RequestRetry defines the retry times when the request failed.
|
// AmpCode contains Amp CLI upstream configuration, management restrictions, and model mappings.
|
||||||
RequestRetry int `yaml:"request-retry" json:"request-retry"`
|
AmpCode AmpCode `yaml:"ampcode" json:"ampcode"`
|
||||||
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
|
|
||||||
MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"`
|
|
||||||
|
|
||||||
// ClaudeKey defines a list of Claude API key configurations as specified in the YAML configuration file.
|
|
||||||
ClaudeKey []ClaudeKey `yaml:"claude-api-key" json:"claude-api-key"`
|
|
||||||
|
|
||||||
// Codex defines a list of Codex API key configurations as specified in the YAML configuration file.
|
|
||||||
CodexKey []CodexKey `yaml:"codex-api-key" json:"codex-api-key"`
|
|
||||||
|
|
||||||
// OpenAICompatibility defines OpenAI API compatibility configurations for external providers.
|
|
||||||
OpenAICompatibility []OpenAICompatibility `yaml:"openai-compatibility" json:"openai-compatibility"`
|
|
||||||
|
|
||||||
// RemoteManagement nests management-related options under 'remote-management'.
|
|
||||||
RemoteManagement RemoteManagement `yaml:"remote-management" json:"-"`
|
|
||||||
|
|
||||||
// Payload defines default and override rules for provider payload parameters.
|
|
||||||
Payload PayloadConfig `yaml:"payload" json:"payload"`
|
|
||||||
|
|
||||||
// OAuthExcludedModels defines per-provider global model exclusions applied to OAuth/file-backed auth entries.
|
// OAuthExcludedModels defines per-provider global model exclusions applied to OAuth/file-backed auth entries.
|
||||||
OAuthExcludedModels map[string][]string `yaml:"oauth-excluded-models,omitempty" json:"oauth-excluded-models,omitempty"`
|
OAuthExcludedModels map[string][]string `yaml:"oauth-excluded-models,omitempty" json:"oauth-excluded-models,omitempty"`
|
||||||
|
|
||||||
|
// Payload defines default and override rules for provider payload parameters.
|
||||||
|
Payload PayloadConfig `yaml:"payload" json:"payload"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// TLSConfig holds HTTPS server settings.
|
// TLSConfig holds HTTPS server settings.
|
||||||
@@ -140,6 +126,26 @@ type AmpModelMapping struct {
|
|||||||
To string `yaml:"to" json:"to"`
|
To string `yaml:"to" json:"to"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// AmpCode groups Amp CLI integration settings including upstream routing,
|
||||||
|
// optional overrides, management route restrictions, and model fallback mappings.
|
||||||
|
type AmpCode struct {
|
||||||
|
// UpstreamURL defines the upstream Amp control plane used for non-provider calls.
|
||||||
|
UpstreamURL string `yaml:"upstream-url" json:"upstream-url"`
|
||||||
|
|
||||||
|
// UpstreamAPIKey optionally overrides the Authorization header when proxying Amp upstream calls.
|
||||||
|
UpstreamAPIKey string `yaml:"upstream-api-key" json:"upstream-api-key"`
|
||||||
|
|
||||||
|
// RestrictManagementToLocalhost restricts Amp management routes (/api/user, /api/threads, etc.)
|
||||||
|
// to only accept connections from localhost (127.0.0.1, ::1). When true, prevents drive-by
|
||||||
|
// browser attacks and remote access to management endpoints. Default: true (recommended).
|
||||||
|
RestrictManagementToLocalhost bool `yaml:"restrict-management-to-localhost" json:"restrict-management-to-localhost"`
|
||||||
|
|
||||||
|
// ModelMappings defines model name mappings for Amp CLI requests.
|
||||||
|
// When Amp requests a model that isn't available locally, these mappings
|
||||||
|
// allow routing to an alternative model that IS available.
|
||||||
|
ModelMappings []AmpModelMapping `yaml:"model-mappings" json:"model-mappings"`
|
||||||
|
}
|
||||||
|
|
||||||
// PayloadConfig defines default and override parameter rules applied to provider payloads.
|
// PayloadConfig defines default and override parameter rules applied to provider payloads.
|
||||||
type PayloadConfig struct {
|
type PayloadConfig struct {
|
||||||
// Default defines rules that only set parameters when they are missing in the payload.
|
// Default defines rules that only set parameters when they are missing in the payload.
|
||||||
@@ -318,7 +324,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
|||||||
cfg.LoggingToFile = false
|
cfg.LoggingToFile = false
|
||||||
cfg.UsageStatisticsEnabled = false
|
cfg.UsageStatisticsEnabled = false
|
||||||
cfg.DisableCooling = false
|
cfg.DisableCooling = false
|
||||||
cfg.AmpRestrictManagementToLocalhost = true // Default to secure: only localhost access
|
cfg.AmpCode.RestrictManagementToLocalhost = true // Default to secure: only localhost access
|
||||||
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
||||||
if optional {
|
if optional {
|
||||||
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
||||||
|
|||||||
Reference in New Issue
Block a user