From bd1678457b7d20a50cb09a5782943e5731c71199 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Wed, 3 Dec 2025 10:42:28 +0800 Subject: [PATCH] refactor(config): consolidate Amp settings into AmpCode struct --- internal/api/modules/amp/amp.go | 16 +++--- internal/api/modules/amp/amp_test.go | 28 ++++++---- internal/config/config.go | 82 +++++++++++++++------------- 3 files changed, 71 insertions(+), 55 deletions(-) diff --git a/internal/api/modules/amp/amp.go b/internal/api/modules/amp/amp.go index 281fda65..fac77bfb 100644 --- a/internal/api/modules/amp/amp.go +++ b/internal/api/modules/amp/amp.go @@ -95,7 +95,8 @@ func (m *AmpModule) Name() string { // This implements the RouteModuleV2 interface with Context. // Routes are registered only once via sync.Once for idempotent behavior. 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) auth := m.getAuthMiddleware(ctx) @@ -104,7 +105,7 @@ func (m *AmpModule) Register(ctx modules.Context) error { var regErr error m.registerOnce.Do(func() { // 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 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 // Cache secrets for 5 minutes to reduce file I/O 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 @@ -136,7 +137,7 @@ func (m *AmpModule) Register(ctx modules.Context) error { // Register management proxy routes (requires upstream) // Restrict to localhost by default for security (prevents drive-by browser attacks) 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.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 { // Update model mappings (hot-reload supported) if m.modelMapper != nil { - log.Infof("amp config updated: reloading %d model mapping(s)", len(cfg.AmpModelMappings)) - m.modelMapper.UpdateMappings(cfg.AmpModelMappings) + settings := cfg.AmpCode + log.Infof("amp config updated: reloading %d model mapping(s)", len(settings.ModelMappings)) + m.modelMapper.UpdateMappings(settings.ModelMappings) } else { 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 } - upstreamURL := strings.TrimSpace(cfg.AmpUpstreamURL) + upstreamURL := strings.TrimSpace(cfg.AmpCode.UpstreamURL) if upstreamURL == "" { log.Warn("Amp upstream URL removed from config, restart required to disable") return nil diff --git a/internal/api/modules/amp/amp_test.go b/internal/api/modules/amp/amp_test.go index 5ae16647..39db6f53 100644 --- a/internal/api/modules/amp/amp_test.go +++ b/internal/api/modules/amp/amp_test.go @@ -56,8 +56,10 @@ func TestAmpModule_Register_WithUpstream(t *testing.T) { m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() }) cfg := &config.Config{ - AmpUpstreamURL: upstream.URL, - AmpUpstreamAPIKey: "test-key", + AmpCode: config.AmpCode{ + UpstreamURL: upstream.URL, + UpstreamAPIKey: "test-key", + }, } 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() }) 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() }} @@ -121,7 +125,9 @@ func TestAmpModule_Register_InvalidUpstream(t *testing.T) { m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() }) 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() }} @@ -151,7 +157,7 @@ func TestAmpModule_OnConfigUpdated_CacheInvalidation(t *testing.T) { } // 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) } @@ -175,7 +181,7 @@ func TestAmpModule_OnConfigUpdated_URLRemoved(t *testing.T) { m.secretSource = ms // 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 { t.Fatalf("unexpected error: %v", err) @@ -187,7 +193,7 @@ func TestAmpModule_OnConfigUpdated_NonMultiSourceSecret(t *testing.T) { m := &AmpModule{enabled: true} 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 if err := m.OnConfigUpdated(cfg); err != nil { @@ -240,8 +246,10 @@ func TestAmpModule_SecretSource_FromConfig(t *testing.T) { // Config with explicit API key cfg := &config.Config{ - AmpUpstreamURL: upstream.URL, - AmpUpstreamAPIKey: "config-key", + AmpCode: config.AmpCode{ + UpstreamURL: upstream.URL, + UpstreamAPIKey: "config-key", + }, } 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() }) - 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() }} if err := m.Register(ctx); err != nil && scenario.configURL != "" { diff --git a/internal/config/config.go b/internal/config/config.go index 16c8b4dc..afe498de 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -26,22 +26,8 @@ type Config struct { // TLS config controls HTTPS server settings. TLS TLSConfig `yaml:"tls" json:"tls"` - // AmpUpstreamURL defines the upstream Amp control plane used for non-provider calls. - AmpUpstreamURL string `yaml:"amp-upstream-url" json:"amp-upstream-url"` - - // 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"` + // RemoteManagement nests management-related options under 'remote-management'. + RemoteManagement RemoteManagement `yaml:"remote-management" json:"-"` // AuthDir is the directory where authentication token files are stored. AuthDir string `yaml:"auth-dir" json:"-"` @@ -58,44 +44,44 @@ type Config struct { // DisableCooling disables quota cooldown scheduling when true. 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 QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"` // WebsocketAuth enables or disables authentication for the WebSocket API. 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 []string `yaml:"generative-language-api-key" json:"generative-language-api-key"` - // GeminiKey defines Gemini API key configurations with optional routing overrides. - GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-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"` + + // 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. // 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"` - // 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"` - - // 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"` + // AmpCode contains Amp CLI upstream configuration, management restrictions, and model mappings. + AmpCode AmpCode `yaml:"ampcode" json:"ampcode"` // 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"` + + // Payload defines default and override rules for provider payload parameters. + Payload PayloadConfig `yaml:"payload" json:"payload"` } // TLSConfig holds HTTPS server settings. @@ -140,6 +126,26 @@ type AmpModelMapping struct { 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. type PayloadConfig struct { // 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.UsageStatisticsEnabled = 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 optional { // In cloud deploy mode, if YAML parsing fails, return empty config instead of error.