diff --git a/core/bifrost.go b/core/bifrost.go index 8288977885..33f7cff1e1 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -959,6 +959,11 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont } } +// GetDropExcessRequests returns the current value of DropExcessRequests +func (bifrost *Bifrost) GetDropExcessRequests() bool { + return bifrost.dropExcessRequests.Load() +} + // UpdateDropExcessRequests updates the DropExcessRequests setting at runtime. // This allows for hot-reloading of this configuration value. func (bifrost *Bifrost) UpdateDropExcessRequests(value bool) { diff --git a/docs/contributing/provider.md b/docs/contributing/provider.md index c8f0166de0..6590284d60 100644 --- a/docs/contributing/provider.md +++ b/docs/contributing/provider.md @@ -80,6 +80,22 @@ type Provider interface { } ``` +### **Meta Configuration Support** + +Some providers require additional configuration beyond API keys. Bifrost supports this through meta configs: + +```go +// In core/schemas/meta/yourprovider.go +type YourProviderMetaConfig struct { + // Add provider-specific fields + Endpoint string `json:"endpoint"` // e.g., Custom API endpoint + Region string `json:"region"` // e.g., Cloud region + ProjectID string `json:"project_id"` // e.g., Cloud project identifier + + // ... other fields (check /core/schemas/provider.go) +} +``` + ### **Provider Structure Template** ```go @@ -488,6 +504,197 @@ func providerRequiresKey(providerKey schemas.ModelProvider) bool { } ``` +## 🌐 **Integration with HTTP Transport** + +The HTTP transport layer requires specific changes to handle provider configuration, meta configs, and model patterns. + +### **1. Provider Recognition** + +Update `transports/bifrost-http/integrations/utils.go`: + +```go +var validProviders = map[schemas.ModelProvider]bool{ + // ... existing providers + schemas.YourProvider: true, // Add this line +} + +// Add model patterns +func isYourProviderModel(model string) bool { + yourProviderPatterns := []string{ + "your-provider-pattern", "your-model-prefix", "yourprovider/", + } + return matchesAnyPattern(model, yourProviderPatterns) +} + +// Add pattern check +func GetProviderFromModel(model string) schemas.ModelProvider { + // ... existing checks + if isYourProviderModel(modelLower) { + return schemas.YourProvider + } +} +``` + +### **2. Meta Configuration Support** + +If your provider needs additional configuration beyond API keys, you'll need to implement meta config support: + +1. **Define Meta Config Structure** (`core/schemas/meta/yourprovider.go`): + +```go +type YourProviderMetaConfig struct { + Endpoint string `json:"endpoint"` // Custom API endpoint + Region string `json:"region"` // Cloud region + ProjectID string `json:"project_id"` // Project identifier +} + +func (c *YourProviderMetaConfig) GetType() string { + return "yourprovider" +} +``` + +2. **Update Store Meta Config Processing** (`transports/bifrost-http/lib/store.go`): + +Add your provider to these three functions: + +```go +// A. Add to parseMetaConfig +func (s *ConfigStore) parseMetaConfig(rawMetaConfig json.RawMessage, provider schemas.ModelProvider) (*schemas.MetaConfig, error) { + switch provider { + // ... existing cases + case schemas.YourProvider: + var config meta.YourProviderMetaConfig + if err := json.Unmarshal(rawMetaConfig, &config); err != nil { + return nil, fmt.Errorf("failed to unmarshal meta config: %w", err) + } + var metaConfig schemas.MetaConfig = &config + return &metaConfig, nil + } + return nil, fmt.Errorf("unsupported provider for meta config: %s", provider) +} + +// B. Add to processMetaConfigEnvVars +func (s *ConfigStore) processMetaConfigEnvVars(rawMetaConfig json.RawMessage, provider schemas.ModelProvider) (json.RawMessage, error) { + switch provider { + // ... existing cases + case schemas.YourProvider: + var config meta.YourProviderMetaConfig + if err := json.Unmarshal(rawMetaConfig, &config); err != nil { + return nil, fmt.Errorf("failed to unmarshal meta config: %w", err) + } + + // Process each field that might contain env vars + endpoint, envVar, err := s.processEnvValue(config.Endpoint) + if err != nil { + return nil, err + } + if envVar != "" { + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.endpoint", provider), + }) + } + config.Endpoint = endpoint + + // Process other fields similarly... + + return json.Marshal(config) + } + return rawMetaConfig, nil +} + +// C. Add to GetProviderConfig for redaction +func (s *ConfigStore) GetProviderConfig(provider schemas.ModelProvider) (*ProviderConfig, error) { + // ... existing code ... + + if configCopy.MetaConfig != nil { + switch m := (*configCopy.MetaConfig).(type) { + // ... existing cases + case *meta.YourProviderMetaConfig: + config := *m + + // Redact or show env vars for each field + path := fmt.Sprintf("providers.%s.meta_config.endpoint", provider) + if envVar, ok := envVarsByPath[path]; ok { + config.Endpoint = "env." + envVar + } else { + config.Endpoint = RedactKey(config.Endpoint) + } + + // Handle other fields... + + var metaConfig schemas.MetaConfig = &config + configCopy.MetaConfig = &metaConfig + } + } + + return &configCopy, nil +} +``` + +### **3. Testing HTTP Transport Integration** + +Add integration tests in `tests/transports-integrations/`: + +```python +# tests/integrations/test_yourprovider.py + +def test_yourprovider_config(): + config = { + "provider": "yourprovider", + "meta_config": { + "endpoint": "env.YOURPROVIDER_ENDPOINT", + "region": "us-east-1" + } + } + # Test config validation + response = client.post("/v1/providers", json=config) + assert response.status_code == 200 + +def test_yourprovider_models(): + # Test model pattern recognition + response = client.post("/v1/chat/completions", json={ + "model": "yourprovider/model-name", + "messages": [{"role": "user", "content": "Hello"}] + }) + assert response.status_code == 200 +``` + +Run the tests: + +```bash +cd tests/transports-integrations +python -m pytest tests/integrations/ -v +``` + +### **4. Configuration Example** + +Document the configuration format for users: + +```json +{ + "providers": { + "yourprovider": { + "keys": [ + { + "value": "env.YOURPROVIDER_API_KEY", + "models": ["*"] + } + ], + "meta_config": { + "endpoint": "env.YOURPROVIDER_ENDPOINT", + "region": "env.YOURPROVIDER_REGION", + "project_id": "env.YOURPROVIDER_PROJECT_ID" + } + } + } +} +``` + +Note: API key handling is automatic - you only need to implement the meta config processing if your provider requires additional configuration beyond API keys. + --- ## 📚 **Documentation Requirements** @@ -539,7 +746,6 @@ result, err := client.ChatCompletionRequest(ctx, &schemas.BifrostRequest{ }, }, }) -``` ## Features @@ -556,7 +762,7 @@ result, err := client.ChatCompletionRequest(ctx, &schemas.BifrostRequest{ | ----------------- | ---------------------- | ---------- | | temperature | temperature | 0.0-2.0 | | max_tokens | max_tokens | Up to 4096 | -```` +``` --- @@ -740,6 +946,7 @@ func (p *YourProviderProvider) convertTools(tools *[]schemas.Tool) []YourProvide return providerTools } ``` +```` --- @@ -759,3 +966,7 @@ func (p *YourProviderProvider) convertTools(tools *[]schemas.Tool) []YourProvide **Ready to build your provider?** 🚀 Check out the existing provider implementations in `core/providers/` for reference, and don't hesitate to ask questions in [GitHub Discussions](https://github.com/maximhq/bifrost/discussions) if you need help! + +``` + +``` diff --git a/transports/bifrost-http/handlers/config.go b/transports/bifrost-http/handlers/config.go index fdfe86d855..3e5342db1a 100644 --- a/transports/bifrost-http/handlers/config.go +++ b/transports/bifrost-http/handlers/config.go @@ -3,11 +3,12 @@ package handlers import ( "encoding/json" "fmt" - "os" + "slices" "github.com/fasthttp/router" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/transports/bifrost-http/lib" "github.com/valyala/fasthttp" ) @@ -16,15 +17,17 @@ import ( type ConfigHandler struct { client *bifrost.Bifrost logger schemas.Logger + store *lib.ConfigStore configPath string } // NewConfigHandler creates a new handler for configuration management. // It requires the Bifrost client, a logger, and the path to the config file to be reloaded. -func NewConfigHandler(client *bifrost.Bifrost, logger schemas.Logger, configPath string) *ConfigHandler { +func NewConfigHandler(client *bifrost.Bifrost, logger schemas.Logger, store *lib.ConfigStore, configPath string) *ConfigHandler { return &ConfigHandler{ client: client, logger: logger, + store: store, configPath: configPath, } } @@ -32,34 +35,66 @@ func NewConfigHandler(client *bifrost.Bifrost, logger schemas.Logger, configPath // RegisterRoutes registers the configuration-related routes. // It adds the `PUT /config` endpoint. func (h *ConfigHandler) RegisterRoutes(r *router.Router) { - r.PUT("/config", h.handleReloadConfig) + r.GET("/config", h.GetConfig) + r.PUT("/config", h.handleUpdateConfig) + r.POST("/config/save", h.SaveConfig) } -// handleReloadConfig re-reads the configuration file and applies updatable settings. +// GetConfig handles GET /config - Get the current configuration +func (h *ConfigHandler) GetConfig(ctx *fasthttp.RequestCtx) { + config := h.store.ClientConfig + SendJSON(ctx, config, h.logger) +} + +// handleUpdateConfig updates the core configuration settings. // Currently, it supports hot-reloading of the `drop_excess_requests` setting. // Note that settings like `prometheus_labels` cannot be changed at runtime. -func (h *ConfigHandler) handleReloadConfig(ctx *fasthttp.RequestCtx) { - var config struct { - BifrostSettings struct { - DropExcessRequests *bool `json:"drop_excess_requests,omitempty"` - } `json:"bifrost_settings"` - } +func (h *ConfigHandler) handleUpdateConfig(ctx *fasthttp.RequestCtx) { + var req lib.ClientConfig - data, err := os.ReadFile(h.configPath) - if err != nil { - SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("failed to read config file: %v", err), h.logger) + if err := json.Unmarshal(ctx.PostBody(), &req); err != nil { + SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("Invalid request format: %v", err), h.logger) return } - if err := json.Unmarshal(data, &config); err != nil { - SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("failed to parse config file: %v", err), h.logger) - return + // Get current config with proper locking + currentConfig := h.store.ClientConfig + updatedConfig := currentConfig + + if req.DropExcessRequests != currentConfig.DropExcessRequests { + h.client.UpdateDropExcessRequests(req.DropExcessRequests) + updatedConfig.DropExcessRequests = req.DropExcessRequests + } + + if !slices.Equal(req.PrometheusLabels, currentConfig.PrometheusLabels) { + updatedConfig.PrometheusLabels = req.PrometheusLabels } - if config.BifrostSettings.DropExcessRequests != nil { - h.client.UpdateDropExcessRequests(*config.BifrostSettings.DropExcessRequests) + if req.InitialPoolSize != currentConfig.InitialPoolSize { + updatedConfig.InitialPoolSize = req.InitialPoolSize } + // Update the store with the new config + h.store.ClientConfig = updatedConfig + ctx.SetStatusCode(fasthttp.StatusOK) - SendJSON(ctx, map[string]interface{}{"status": "config reloaded", "drop_excess_requests": config.BifrostSettings.DropExcessRequests}, h.logger) + SendJSON(ctx, map[string]any{ + "status": "success", + "message": "Configuration updated successfully", + }, h.logger) +} + +// SaveConfig handles POST /config/save - Persist current configuration to JSON file +func (h *ConfigHandler) SaveConfig(ctx *fasthttp.RequestCtx) { + // Save current configuration back to the original JSON file + if err := h.store.SaveConfig(); err != nil { + h.logger.Warn(fmt.Sprintf("Failed to save configuration: %v", err)) + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to save configuration: %v", err), h.logger) + return + } + + SendJSON(ctx, map[string]any{ + "status": "success", + "message": "Configuration saved successfully", + }, h.logger) } diff --git a/transports/bifrost-http/handlers/mcp.go b/transports/bifrost-http/handlers/mcp.go index 7a952cc23a..dd6bb54db3 100644 --- a/transports/bifrost-http/handlers/mcp.go +++ b/transports/bifrost-http/handlers/mcp.go @@ -69,18 +69,13 @@ func (h *MCPHandler) ExecuteTool(ctx *fasthttp.RequestCtx) { } // Send successful response - ctx.SetStatusCode(fasthttp.StatusOK) - ctx.SetContentType("application/json") - if encodeErr := json.NewEncoder(ctx).Encode(resp); encodeErr != nil { - h.logger.Warn(fmt.Sprintf("Failed to encode response: %v", encodeErr)) - SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to encode response: %v", encodeErr), h.logger) - } + SendJSON(ctx, resp, h.logger) } // GetMCPClients handles GET /mcp/clients - Get all MCP clients func (h *MCPHandler) GetMCPClients(ctx *fasthttp.RequestCtx) { // Get clients from store config - configsInStore := h.store.GetMCPConfig() + configsInStore := h.store.MCPConfig if configsInStore == nil { SendJSON(ctx, []schemas.MCPClient{}, h.logger) return @@ -105,12 +100,17 @@ func (h *MCPHandler) GetMCPClients(ctx *fasthttp.RequestCtx) { for _, configClient := range configsInStore.ClientConfigs { if connectedClient, exists := connectedClientsMap[configClient.Name]; exists { // Client is connected, use the actual client data - clients = append(clients, connectedClient) + clients = append(clients, schemas.MCPClient{ + Name: connectedClient.Name, + Config: h.store.RedactMCPClientConfig(connectedClient.Config), + Tools: connectedClient.Tools, + State: connectedClient.State, + }) } else { // Client is in config but not connected, mark as errored clients = append(clients, schemas.MCPClient{ Name: configClient.Name, - Config: configClient, + Config: h.store.RedactMCPClientConfig(configClient), Tools: []string{}, // No tools available since connection failed State: schemas.MCPConnectionStateError, }) diff --git a/transports/bifrost-http/handlers/providers.go b/transports/bifrost-http/handlers/providers.go index bdf685eb21..7f0f0497dc 100644 --- a/transports/bifrost-http/handlers/providers.go +++ b/transports/bifrost-http/handlers/providers.go @@ -5,6 +5,7 @@ package handlers import ( "encoding/json" "fmt" + "sort" "github.com/fasthttp/router" bifrost "github.com/maximhq/bifrost/core" @@ -79,9 +80,6 @@ func (h *ProviderHandler) RegisterRoutes(r *router.Router) { r.POST("/providers", h.AddProvider) r.PUT("/providers/{provider}", h.UpdateProvider) r.DELETE("/providers/{provider}", h.DeleteProvider) - - // Configuration persistence - r.POST("/config/save", h.SaveConfig) } // ListProviders handles GET /providers - List all providers @@ -93,8 +91,14 @@ func (h *ProviderHandler) ListProviders(ctx *fasthttp.RequestCtx) { } var providerResponses []ProviderResponse + + // Sort providers alphabetically + sort.Slice(providers, func(i, j int) bool { + return string(providers[i]) < string(providers[j]) + }) + for _, provider := range providers { - config, err := h.store.GetProviderConfig(provider) + config, err := h.store.GetProviderConfigRedacted(provider) if err != nil { h.logger.Warn(fmt.Sprintf("Failed to get config for provider %s: %v", provider, err)) // Include provider even if config fetch fails @@ -123,7 +127,7 @@ func (h *ProviderHandler) GetProvider(ctx *fasthttp.RequestCtx) { return } - config, err := h.store.GetProviderConfig(provider) + config, err := h.store.GetProviderConfigRedacted(provider) if err != nil { SendError(ctx, fasthttp.StatusNotFound, fmt.Sprintf("Provider not found: %v", err), h.logger) return @@ -166,7 +170,7 @@ func (h *ProviderHandler) AddProvider(ctx *fasthttp.RequestCtx) { } // Check if provider already exists - if _, err := h.store.GetProviderConfig(req.Provider); err == nil { + if _, err := h.store.GetProviderConfigRedacted(req.Provider); err == nil { SendError(ctx, fasthttp.StatusConflict, fmt.Sprintf("Provider %s already exists", req.Provider), h.logger) return } @@ -220,8 +224,8 @@ func (h *ProviderHandler) UpdateProvider(ctx *fasthttp.RequestCtx) { return } - // Check if provider exists - oldConfig, err := h.store.GetProviderConfig(provider) + // Get the raw config to access actual values for merging with redacted request values + oldConfigRaw, err := h.store.GetProviderConfigRaw(provider) if err != nil { SendError(ctx, fasthttp.StatusNotFound, fmt.Sprintf("Provider not found: %v", err), h.logger) return @@ -229,24 +233,81 @@ func (h *ProviderHandler) UpdateProvider(ctx *fasthttp.RequestCtx) { // Construct ProviderConfig from individual fields config := lib.ProviderConfig{ - Keys: oldConfig.Keys, - NetworkConfig: oldConfig.NetworkConfig, - ConcurrencyAndBufferSize: oldConfig.ConcurrencyAndBufferSize, + Keys: oldConfigRaw.Keys, + NetworkConfig: oldConfigRaw.NetworkConfig, + ConcurrencyAndBufferSize: oldConfigRaw.ConcurrencyAndBufferSize, } - // Validate required keys (at least one key must be provided) + // For now, don't replace any environment keys - preserve all existing ones + // TODO: Implement proper tracking of which env keys should be dropped + envKeysToReplace := make(map[string]struct{}) + + // Validate and process keys if req.Keys != nil { if len(req.Keys) == 0 && provider != schemas.Vertex && provider != schemas.Ollama { SendError(ctx, fasthttp.StatusBadRequest, "At least one API key is required", h.logger) return } + + // Create a map of old keys by model patterns for quick lookup + oldKeysByModels := make(map[string][]schemas.Key) + for _, oldKey := range oldConfigRaw.Keys { + for _, model := range oldKey.Models { + oldKeysByModels[model] = append(oldKeysByModels[model], oldKey) + } + } + + // Process each key in the request + for i, newKey := range req.Keys { + // If the key is redacted, try to find and use the old key for the same models + if lib.IsRedacted(newKey.Value) { + // Look for matching old keys + var matchingKeys []schemas.Key + for _, model := range newKey.Models { + if oldKeys, exists := oldKeysByModels[model]; exists { + matchingKeys = append(matchingKeys, oldKeys...) + } + } + + // If we found matching keys, use the most appropriate one + if len(matchingKeys) > 0 { + // Try to find a key that matches all the same models + var bestMatch schemas.Key + bestMatchScore := 0 + + for _, oldKey := range matchingKeys { + // Calculate how many models match between the old and new key + matchCount := 0 + oldModelsMap := make(map[string]bool) + for _, m := range oldKey.Models { + oldModelsMap[m] = true + } + + for _, m := range newKey.Models { + if oldModelsMap[m] { + matchCount++ + } + } + + // Update best match if this key has more matching models + if matchCount > bestMatchScore { + bestMatch = oldKey + bestMatchScore = matchCount + } + } + + // Use the best matching key's value + req.Keys[i].Value = bestMatch.Value + } + } + } config.Keys = req.Keys } // Handle meta config if provided if req.MetaConfig != nil && len(*req.MetaConfig) > 0 { - // Convert to appropriate meta config type based on provider - metaConfig, err := h.convertToProviderMetaConfig(provider, *req.MetaConfig) + // Merge new meta config with old, preserving redacted values + metaConfig, err := h.mergeMetaConfig(provider, oldConfigRaw.MetaConfig, *req.MetaConfig) if err != nil { SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("Invalid meta config: %v", err), h.logger) return @@ -273,14 +334,14 @@ func (h *ProviderHandler) UpdateProvider(ctx *fasthttp.RequestCtx) { config.ProxyConfig = req.ProxyConfig // Update provider config in store (env vars will be processed by store) - if err := h.store.UpdateProviderConfig(provider, config); err != nil { + if err := h.store.UpdateProviderConfig(provider, config, envKeysToReplace); err != nil { h.logger.Warn(fmt.Sprintf("Failed to update provider %s: %v", provider, err)) SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to update provider: %v", err), h.logger) return } - if config.ConcurrencyAndBufferSize.Concurrency != oldConfig.ConcurrencyAndBufferSize.Concurrency || - config.ConcurrencyAndBufferSize.BufferSize != oldConfig.ConcurrencyAndBufferSize.BufferSize { + if config.ConcurrencyAndBufferSize.Concurrency != oldConfigRaw.ConcurrencyAndBufferSize.Concurrency || + config.ConcurrencyAndBufferSize.BufferSize != oldConfigRaw.ConcurrencyAndBufferSize.BufferSize { // Update concurrency and queue configuration in Bifrost if err := h.client.UpdateProviderConcurrency(provider); err != nil { // Note: Store update succeeded, continue but log the concurrency update failure @@ -302,7 +363,7 @@ func (h *ProviderHandler) DeleteProvider(ctx *fasthttp.RequestCtx) { } // Check if provider exists - if _, err := h.store.GetProviderConfig(provider); err != nil { + if _, err := h.store.GetProviderConfigRedacted(provider); err != nil { SendError(ctx, fasthttp.StatusNotFound, fmt.Sprintf("Provider not found: %v", err), h.logger) return } @@ -323,25 +384,6 @@ func (h *ProviderHandler) DeleteProvider(ctx *fasthttp.RequestCtx) { SendJSON(ctx, response, h.logger) } -// SaveConfig handles POST /config/save - Persist current configuration to JSON file -func (h *ProviderHandler) SaveConfig(ctx *fasthttp.RequestCtx) { - // Save current configuration back to the original JSON file - if err := h.store.SaveConfig(); err != nil { - h.logger.Warn(fmt.Sprintf("Failed to save configuration: %v", err)) - SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to save configuration: %v", err), h.logger) - return - } - - h.logger.Info("Configuration saved successfully") - - response := map[string]interface{}{ - "status": "success", - "message": "Configuration saved successfully", - } - - SendJSON(ctx, response, h.logger) -} - // convertToProviderMetaConfig converts a generic map to the appropriate provider-specific meta config func (h *ProviderHandler) convertToProviderMetaConfig(provider schemas.ModelProvider, metaConfigMap map[string]interface{}) (*schemas.MetaConfig, error) { if len(metaConfigMap) == 0 { @@ -385,6 +427,96 @@ func (h *ProviderHandler) convertToProviderMetaConfig(provider schemas.ModelProv } } +// mergeMetaConfig merges new meta config with old, preserving values that are redacted in the new config +func (h *ProviderHandler) mergeMetaConfig(provider schemas.ModelProvider, oldConfig *schemas.MetaConfig, newConfigMap map[string]interface{}) (*schemas.MetaConfig, error) { + if oldConfig == nil || len(newConfigMap) == 0 { + return h.convertToProviderMetaConfig(provider, newConfigMap) + } + + switch provider { + case schemas.Azure: + var newAzureConfig meta.AzureMetaConfig + newConfigJSON, _ := json.Marshal(newConfigMap) + if err := json.Unmarshal(newConfigJSON, &newAzureConfig); err != nil { + return nil, fmt.Errorf("failed to unmarshal new Azure meta config: %w", err) + } + + oldAzureConfig, ok := (*oldConfig).(*meta.AzureMetaConfig) + if !ok { + return nil, fmt.Errorf("existing meta config type mismatch: expected AzureMetaConfig") + } + + // Preserve old values if new ones are redacted + if lib.IsRedacted(newAzureConfig.Endpoint) { + newAzureConfig.Endpoint = oldAzureConfig.Endpoint + } + if newAzureConfig.APIVersion != nil && oldAzureConfig.APIVersion != nil && lib.IsRedacted(*newAzureConfig.APIVersion) { + newAzureConfig.APIVersion = oldAzureConfig.APIVersion + } + + var metaConfig schemas.MetaConfig = &newAzureConfig + return &metaConfig, nil + + case schemas.Bedrock: + var newBedrockConfig meta.BedrockMetaConfig + newConfigJSON, _ := json.Marshal(newConfigMap) + if err := json.Unmarshal(newConfigJSON, &newBedrockConfig); err != nil { + return nil, fmt.Errorf("failed to unmarshal new Bedrock meta config: %w", err) + } + + oldBedrockConfig, ok := (*oldConfig).(*meta.BedrockMetaConfig) + if !ok { + return nil, fmt.Errorf("existing meta config type mismatch: expected BedrockMetaConfig") + } + + // Preserve old values if new ones are redacted + if lib.IsRedacted(newBedrockConfig.SecretAccessKey) { + newBedrockConfig.SecretAccessKey = oldBedrockConfig.SecretAccessKey + } + if newBedrockConfig.Region != nil && oldBedrockConfig.Region != nil && lib.IsRedacted(*newBedrockConfig.Region) { + newBedrockConfig.Region = oldBedrockConfig.Region + } + if newBedrockConfig.SessionToken != nil && oldBedrockConfig.SessionToken != nil && lib.IsRedacted(*newBedrockConfig.SessionToken) { + newBedrockConfig.SessionToken = oldBedrockConfig.SessionToken + } + if newBedrockConfig.ARN != nil && oldBedrockConfig.ARN != nil && lib.IsRedacted(*newBedrockConfig.ARN) { + newBedrockConfig.ARN = oldBedrockConfig.ARN + } + + var metaConfig schemas.MetaConfig = &newBedrockConfig + return &metaConfig, nil + + case schemas.Vertex: + var newVertexConfig meta.VertexMetaConfig + newConfigJSON, _ := json.Marshal(newConfigMap) + if err := json.Unmarshal(newConfigJSON, &newVertexConfig); err != nil { + return nil, fmt.Errorf("failed to unmarshal new Vertex meta config: %w", err) + } + + oldVertexConfig, ok := (*oldConfig).(*meta.VertexMetaConfig) + if !ok { + return nil, fmt.Errorf("existing meta config type mismatch: expected VertexMetaConfig") + } + + // Preserve old values if new ones are redacted + if lib.IsRedacted(newVertexConfig.ProjectID) { + newVertexConfig.ProjectID = oldVertexConfig.ProjectID + } + if lib.IsRedacted(newVertexConfig.Region) { + newVertexConfig.Region = oldVertexConfig.Region + } + if lib.IsRedacted(newVertexConfig.AuthCredentials) { + newVertexConfig.AuthCredentials = oldVertexConfig.AuthCredentials + } + + var metaConfig schemas.MetaConfig = &newVertexConfig + return &metaConfig, nil + + default: + return nil, nil + } +} + func (h *ProviderHandler) getProviderResponseFromConfig(provider schemas.ModelProvider, config lib.ProviderConfig) ProviderResponse { if config.NetworkConfig == nil { config.NetworkConfig = &schemas.DefaultNetworkConfig diff --git a/transports/bifrost-http/lib/account.go b/transports/bifrost-http/lib/account.go index 757b7bda81..0e55470aa4 100644 --- a/transports/bifrost-http/lib/account.go +++ b/transports/bifrost-http/lib/account.go @@ -40,7 +40,7 @@ func (baseAccount *BaseAccount) GetKeysForProvider(providerKey schemas.ModelProv return nil, fmt.Errorf("store not initialized") } - config, err := baseAccount.store.GetProviderConfig(providerKey) + config, err := baseAccount.store.GetProviderConfigRaw(providerKey) if err != nil { return nil, err } @@ -56,7 +56,7 @@ func (baseAccount *BaseAccount) GetConfigForProvider(providerKey schemas.ModelPr return nil, fmt.Errorf("store not initialized") } - config, err := baseAccount.store.GetProviderConfig(providerKey) + config, err := baseAccount.store.GetProviderConfigRaw(providerKey) if err != nil { return nil, err } diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go index 733a500a31..80d0fe000a 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -3,16 +3,17 @@ package lib import ( - "encoding/json" - "fmt" - "log" - "os" - "strings" - "github.com/maximhq/bifrost/core/schemas" - "github.com/maximhq/bifrost/core/schemas/meta" ) +// ClientConfig represents the core configuration for Bifrost HTTP transport and the Bifrost Client. +// It includes settings for excess request handling, Prometheus metrics, and initial pool size. +type ClientConfig struct { + DropExcessRequests bool `json:"drop_excess_requests"` + PrometheusLabels []string `json:"prometheus_labels"` + InitialPoolSize int `json:"initial_pool_size"` +} + // ProviderConfig represents the configuration for a specific AI model provider. // It includes API keys, network settings, provider-specific metadata, and concurrency settings. type ProviderConfig struct { @@ -29,152 +30,7 @@ type ConfigMap map[schemas.ModelProvider]ProviderConfig // BifrostHTTPConfig represents the complete configuration structure for Bifrost HTTP transport. // It includes both provider configurations and MCP configuration. type BifrostHTTPConfig struct { + ClientConfig *ClientConfig `json:"client"` // Client configuration ProviderConfig ConfigMap `json:"providers"` // Provider configurations MCPConfig *schemas.MCPConfig `json:"mcp"` // MCP configuration (optional) } - -// ReadMCPKeys reads environment variables from the environment and updates the MCP configurations. -// It replaces values starting with "env." in the connection_string field with actual values from the environment. -// Returns an error if any required environment variable is missing. -func (config *BifrostHTTPConfig) ReadMCPKeys() error { - if config.MCPConfig == nil { - return nil // No MCP config to process - } - - // Helper function to check and replace env values - replaceEnvValue := func(value string) (string, error) { - if strings.HasPrefix(value, "env.") { - envKey := strings.TrimPrefix(value, "env.") - if envValue := os.Getenv(envKey); envValue != "" { - return envValue, nil - } - return "", fmt.Errorf("environment variable %s not found in the environment", envKey) - } - return value, nil - } - - // Process each client config - for i, clientConfig := range config.MCPConfig.ClientConfigs { - // Process ConnectionString if present - if clientConfig.ConnectionString != nil { - newValue, err := replaceEnvValue(*clientConfig.ConnectionString) - if err != nil { - return fmt.Errorf("MCP client %s: %w", clientConfig.Name, err) - } - config.MCPConfig.ClientConfigs[i].ConnectionString = &newValue - } - } - - return nil -} - -// readConfig reads and parses the configuration file. -// It handles case conversion for provider names and sets up provider-specific metadata. -// Returns a BifrostHTTPConfig containing both provider and MCP configurations. -// Panics if the config file cannot be read or parsed. -// -// In the config file, use placeholder keys (e.g., env.OPENAI_API_KEY) instead of hardcoding actual values. -// These placeholders will be replaced with the corresponding values from the environment variables. -// Example: -// -// "providers": { -// "openAI": { -// "keys":[{ -// "value": "env.OPENAI_API_KEY" -// "models": ["gpt-4o-mini", "gpt-4-turbo"], -// "weight": 1.0 -// }] -// } -// }, -// "mcp": { -// "client_configs": [...] -// } -// -// In this example, OPENAI_API_KEY refers to a key in the environment variables. At runtime, its value will be used to replace the placeholder. -// Same setup applies to keys in meta configs of all the providers. -// Example: -// -// "meta_config": { -// "secret_access_key": "env.AWS_SECRET_ACCESS_KEY" -// "region": "env.AWS_REGION" -// } -// -// In this example, AWS_SECRET_ACCESS_KEY and AWS_REGION refer to keys in environment variables. -func ReadConfig(configLocation string) *BifrostHTTPConfig { - data, err := os.ReadFile(configLocation) - if err != nil { - log.Fatalf("failed to read config JSON file: %v", err) - } - - // First unmarshal into the new structure - var fullConfig BifrostHTTPConfig - if err := json.Unmarshal(data, &fullConfig); err != nil { - log.Fatalf("failed to unmarshal JSON: %v", err) - } - - if fullConfig.ProviderConfig == nil { - log.Fatalf("providers section is required in config") - } - - // Process provider configurations - convert string keys to lowercase provider names and handle meta configs - processedProviders := make(ConfigMap) - - // First unmarshal providers into a map with string keys to handle case conversion - var rawProviders map[string]ProviderConfig - if providersBytes, err := json.Marshal(fullConfig.ProviderConfig); err != nil { - log.Fatalf("failed to marshal providers: %v", err) - } else if err := json.Unmarshal(providersBytes, &rawProviders); err != nil { - log.Fatalf("failed to unmarshal providers: %v", err) - } - - // Create a temporary structure to unmarshal the full JSON with proper meta configs - var tempConfig struct { - Providers map[string]struct { - MetaConfig json.RawMessage `json:"meta_config"` - } `json:"providers"` - } - - if err := json.Unmarshal(data, &tempConfig); err != nil { - log.Fatalf("failed to unmarshal configuration file: %v\n\n Please check your configuration file for proper JSON formatting and meta_config structure", err) - } else { - for rawProvider, cfg := range rawProviders { - provider := schemas.ModelProvider(strings.ToLower(rawProvider)) - - // Get the raw meta config for this provider - if tempProvider, exists := tempConfig.Providers[rawProvider]; exists && len(tempProvider.MetaConfig) > 0 { - switch provider { - case schemas.Azure: - var azureMetaConfig meta.AzureMetaConfig - if err := json.Unmarshal(tempProvider.MetaConfig, &azureMetaConfig); err != nil { - log.Printf("warning: failed to unmarshal Azure meta config: %v", err) - } else { - var metaConfig schemas.MetaConfig = &azureMetaConfig - cfg.MetaConfig = &metaConfig - } - case schemas.Bedrock: - var bedrockMetaConfig meta.BedrockMetaConfig - if err := json.Unmarshal(tempProvider.MetaConfig, &bedrockMetaConfig); err != nil { - log.Printf("warning: failed to unmarshal Bedrock meta config: %v", err) - } else { - var metaConfig schemas.MetaConfig = &bedrockMetaConfig - cfg.MetaConfig = &metaConfig - } - case schemas.Vertex: - var vertexMetaConfig meta.VertexMetaConfig - if err := json.Unmarshal(tempProvider.MetaConfig, &vertexMetaConfig); err != nil { - log.Printf("warning: failed to unmarshal Vertex meta config: %v", err) - } else { - var metaConfig schemas.MetaConfig = &vertexMetaConfig - cfg.MetaConfig = &metaConfig - } - } - } - - processedProviders[provider] = cfg - } - - } - - fullConfig.ProviderConfig = processedProviders - return &fullConfig -} diff --git a/transports/bifrost-http/lib/store.go b/transports/bifrost-http/lib/store.go index bcc3f29a88..ac27a14126 100644 --- a/transports/bifrost-http/lib/store.go +++ b/transports/bifrost-http/lib/store.go @@ -33,15 +33,28 @@ type ConfigStore struct { client *bifrost.Bifrost // In-memory storage - providers map[schemas.ModelProvider]ProviderConfig - mcpConfig *schemas.MCPConfig + ClientConfig ClientConfig + Providers map[schemas.ModelProvider]ProviderConfig + MCPConfig *schemas.MCPConfig + + // Track which keys come from environment variables + EnvKeys map[string][]EnvKeyInfo +} + +// EnvKeyInfo stores information about a key sourced from environment +type EnvKeyInfo struct { + EnvVar string // The environment variable name (without env. prefix) + Provider string // The provider this key belongs to (empty for core/mcp configs) + KeyType string // Type of key (e.g., "api_key", "meta_config", "connection_string") + ConfigPath string // Path in config where this env var is used } // NewConfigStore creates a new in-memory configuration store instance. func NewConfigStore(logger schemas.Logger) (*ConfigStore, error) { return &ConfigStore{ logger: logger, - providers: make(map[schemas.ModelProvider]ProviderConfig), + Providers: make(map[schemas.ModelProvider]ProviderConfig), + EnvKeys: make(map[string][]EnvKeyInfo), }, nil } @@ -70,6 +83,7 @@ func (s *ConfigStore) LoadFromConfig(configPath string) error { // Parse the JSON directly var configData struct { + Client json.RawMessage `json:"client"` Providers map[string]json.RawMessage `json:"providers"` MCP json.RawMessage `json:"mcp,omitempty"` } @@ -78,6 +92,15 @@ func (s *ConfigStore) LoadFromConfig(configPath string) error { return fmt.Errorf("failed to unmarshal config: %w", err) } + // Process core configuration if present + if len(configData.Client) > 0 { + var clientConfig ClientConfig + if err := json.Unmarshal(configData.Client, &clientConfig); err != nil { + return fmt.Errorf("failed to unmarshal client config: %w", err) + } + s.ClientConfig = clientConfig + } + // Process provider configurations processedProviders := make(map[schemas.ModelProvider]ProviderConfig) @@ -102,12 +125,16 @@ func (s *ConfigStore) LoadFromConfig(configPath string) error { // Process each provider configuration for rawProviderName, cfg := range rawProviders { + newEnvKeys := make(map[string]struct{}) + provider := schemas.ModelProvider(strings.ToLower(rawProviderName)) // Process meta config if it exists if tempProvider, exists := tempConfig.Providers[rawProviderName]; exists && len(tempProvider.MetaConfig) > 0 { - processedMetaConfig, err := s.processMetaConfigEnvVars(tempProvider.MetaConfig, provider) + processedMetaConfig, envKeys, err := s.processMetaConfigEnvVars(tempProvider.MetaConfig, provider) + if err != nil { + s.cleanupEnvKeys(string(provider), "", envKeys) s.logger.Warn(fmt.Sprintf("failed to process env vars in meta config for %s: %v", provider, err)) continue } @@ -115,7 +142,8 @@ func (s *ConfigStore) LoadFromConfig(configPath string) error { // Parse and set the meta config metaConfig, err := s.parseMetaConfig(processedMetaConfig, provider) if err != nil { - s.logger.Warn(fmt.Sprintf("failed to parse meta config for %s: %v", provider, err)) + s.cleanupEnvKeys(string(provider), "", envKeys) + s.logger.Warn(fmt.Sprintf("failed to process meta config for %s: %v", provider, err)) continue } else { cfg.MetaConfig = metaConfig @@ -124,19 +152,31 @@ func (s *ConfigStore) LoadFromConfig(configPath string) error { // Process environment variables in keys for i, key := range cfg.Keys { - processedValue, err := s.replaceEnvValue(key.Value) + processedValue, envVar, err := s.processEnvValue(key.Value) if err != nil { + s.cleanupEnvKeys(string(provider), "", newEnvKeys) s.logger.Warn(fmt.Sprintf("failed to process env vars in keys for %s: %v", provider, err)) continue } cfg.Keys[i].Value = processedValue + + // Track environment key if it came from env + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "api_key", + ConfigPath: fmt.Sprintf("providers.%s.keys[%d]", provider, i), + }) + } } processedProviders[provider] = cfg } // Store processed configurations in memory - s.providers = processedProviders + s.Providers = processedProviders // Parse MCP config if present if len(configData.MCP) > 0 { @@ -145,7 +185,7 @@ func (s *ConfigStore) LoadFromConfig(configPath string) error { s.logger.Warn(fmt.Sprintf("failed to parse MCP config: %v", err)) } else { // Process environment variables in MCP config - s.mcpConfig = &mcpConfig + s.MCPConfig = &mcpConfig s.processMCPEnvVars() } } @@ -154,6 +194,25 @@ func (s *ConfigStore) LoadFromConfig(configPath string) error { return nil } +// processEnvValue checks and replaces environment variable references in configuration values. +// Returns the processed value and the environment variable name if it was an env reference. +// Supports the "env.VARIABLE_NAME" syntax for referencing environment variables. +// This enables secure configuration management without hardcoding sensitive values. +// +// Examples: +// - "env.OPENAI_API_KEY" -> actual value from OPENAI_API_KEY environment variable +// - "sk-1234567890" -> returned as-is (no env prefix) +func (s *ConfigStore) processEnvValue(value string) (string, string, error) { + if strings.HasPrefix(value, "env.") { + envKey := strings.TrimPrefix(value, "env.") + if envValue := os.Getenv(envKey); envValue != "" { + return envValue, envKey, nil + } + return "", envKey, fmt.Errorf("environment variable %s not found", envKey) + } + return value, "", nil +} + // WriteConfigToFile writes the current in-memory configuration back to a JSON file // in the exact same format that LoadFromConfig expects. This enables persistence // of runtime configuration changes. @@ -167,13 +226,15 @@ func (s *ConfigStore) WriteConfigToFile(configPath string) error { output := struct { Providers map[string]interface{} `json:"providers"` MCP *schemas.MCPConfig `json:"mcp,omitempty"` + Client ClientConfig `json:"client,omitempty"` }{ Providers: make(map[string]interface{}), - MCP: s.mcpConfig, + MCP: s.MCPConfig, + Client: s.ClientConfig, } // Convert providers back to the original format - for provider, config := range s.providers { + for provider, config := range s.Providers { providerName := string(provider) // Create provider config without processed values (keep env.* references) @@ -250,26 +311,6 @@ func (s *ConfigStore) parseMetaConfig(rawMetaConfig json.RawMessage, provider sc return nil, fmt.Errorf("unsupported provider for meta config: %s", provider) } -// replaceEnvValue checks and replaces environment variable references in configuration values. -// Supports the "env.VARIABLE_NAME" syntax for referencing environment variables. -// This enables secure configuration management without hardcoding sensitive values. -// -// Examples: -// - "env.OPENAI_API_KEY" -> actual value from OPENAI_API_KEY environment variable -// - "sk-1234567890" -> returned as-is (no env prefix) -// -// Returns an error if the referenced environment variable is not found. -func (s *ConfigStore) replaceEnvValue(value string) (string, error) { - if strings.HasPrefix(value, "env.") { - envKey := strings.TrimPrefix(value, "env.") - if envValue := os.Getenv(envKey); envValue != "" { - return envValue, nil - } - return "", fmt.Errorf("environment variable %s not found", envKey) - } - return value, nil -} - // processMetaConfigEnvVars processes environment variables in provider-specific meta configurations. // This method handles the provider-specific meta config structures and processes environment // variables in their fields, ensuring type safety and proper field handling. @@ -281,111 +322,196 @@ func (s *ConfigStore) replaceEnvValue(value string) (string, error) { // // For unsupported providers, the meta config is returned unchanged. // This approach ensures type safety while supporting environment variable substitution. -func (s *ConfigStore) processMetaConfigEnvVars(rawMetaConfig json.RawMessage, provider schemas.ModelProvider) (json.RawMessage, error) { +func (s *ConfigStore) processMetaConfigEnvVars(rawMetaConfig json.RawMessage, provider schemas.ModelProvider) (json.RawMessage, map[string]struct{}, error) { + // Track new environment variables + newEnvKeys := make(map[string]struct{}) + switch provider { case schemas.Azure: var azureMetaConfig meta.AzureMetaConfig if err := json.Unmarshal(rawMetaConfig, &azureMetaConfig); err != nil { - return nil, fmt.Errorf("failed to unmarshal Azure meta config: %w", err) + return nil, newEnvKeys, fmt.Errorf("failed to unmarshal Azure meta config: %w", err) } - endpoint, err := s.replaceEnvValue(azureMetaConfig.Endpoint) + endpoint, envVar, err := s.processEnvValue(azureMetaConfig.Endpoint) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.endpoint", provider), + }) } azureMetaConfig.Endpoint = endpoint + if azureMetaConfig.APIVersion != nil { - apiVersion, err := s.replaceEnvValue(*azureMetaConfig.APIVersion) + apiVersion, envVar, err := s.processEnvValue(*azureMetaConfig.APIVersion) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.api_version", provider), + }) } azureMetaConfig.APIVersion = &apiVersion } processedJSON, err := json.Marshal(azureMetaConfig) if err != nil { - return nil, fmt.Errorf("failed to marshal processed Azure meta config: %w", err) + return nil, newEnvKeys, fmt.Errorf("failed to marshal processed Azure meta config: %w", err) } - return processedJSON, nil + return processedJSON, newEnvKeys, nil case schemas.Bedrock: var bedrockMetaConfig meta.BedrockMetaConfig if err := json.Unmarshal(rawMetaConfig, &bedrockMetaConfig); err != nil { - return nil, fmt.Errorf("failed to unmarshal Bedrock meta config: %w", err) + return nil, newEnvKeys, fmt.Errorf("failed to unmarshal Bedrock meta config: %w", err) } - secretAccessKey, err := s.replaceEnvValue(bedrockMetaConfig.SecretAccessKey) + secretAccessKey, envVar, err := s.processEnvValue(bedrockMetaConfig.SecretAccessKey) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.secret_access_key", provider), + }) } bedrockMetaConfig.SecretAccessKey = secretAccessKey if bedrockMetaConfig.Region != nil { - region, err := s.replaceEnvValue(*bedrockMetaConfig.Region) + region, envVar, err := s.processEnvValue(*bedrockMetaConfig.Region) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.region", provider), + }) } bedrockMetaConfig.Region = ®ion } if bedrockMetaConfig.SessionToken != nil { - sessionToken, err := s.replaceEnvValue(*bedrockMetaConfig.SessionToken) + sessionToken, envVar, err := s.processEnvValue(*bedrockMetaConfig.SessionToken) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.session_token", provider), + }) } bedrockMetaConfig.SessionToken = &sessionToken } if bedrockMetaConfig.ARN != nil { - arn, err := s.replaceEnvValue(*bedrockMetaConfig.ARN) + arn, envVar, err := s.processEnvValue(*bedrockMetaConfig.ARN) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.arn", provider), + }) } bedrockMetaConfig.ARN = &arn } processedJSON, err := json.Marshal(bedrockMetaConfig) if err != nil { - return nil, fmt.Errorf("failed to marshal processed Bedrock meta config: %w", err) + return nil, newEnvKeys, fmt.Errorf("failed to marshal processed Bedrock meta config: %w", err) } - return processedJSON, nil + return processedJSON, newEnvKeys, nil case schemas.Vertex: var vertexMetaConfig meta.VertexMetaConfig if err := json.Unmarshal(rawMetaConfig, &vertexMetaConfig); err != nil { - return nil, fmt.Errorf("failed to unmarshal Vertex meta config: %w", err) + return nil, newEnvKeys, fmt.Errorf("failed to unmarshal Vertex meta config: %w", err) } - projectID, err := s.replaceEnvValue(vertexMetaConfig.ProjectID) + projectID, envVar, err := s.processEnvValue(vertexMetaConfig.ProjectID) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.project_id", provider), + }) } vertexMetaConfig.ProjectID = projectID - region, err := s.replaceEnvValue(vertexMetaConfig.Region) + region, envVar, err := s.processEnvValue(vertexMetaConfig.Region) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.region", provider), + }) } vertexMetaConfig.Region = region - authCredentials, err := s.replaceEnvValue(vertexMetaConfig.AuthCredentials) + authCredentials, envVar, err := s.processEnvValue(vertexMetaConfig.AuthCredentials) if err != nil { - return nil, err + return nil, newEnvKeys, err + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "meta_config", + ConfigPath: fmt.Sprintf("providers.%s.meta_config.auth_credentials", provider), + }) } vertexMetaConfig.AuthCredentials = authCredentials processedJSON, err := json.Marshal(vertexMetaConfig) if err != nil { - return nil, fmt.Errorf("failed to marshal processed Vertex meta config: %w", err) + return nil, newEnvKeys, fmt.Errorf("failed to marshal processed Vertex meta config: %w", err) } - return processedJSON, nil + return processedJSON, newEnvKeys, nil } - return rawMetaConfig, nil + return rawMetaConfig, newEnvKeys, nil } -// GetProviderConfig retrieves a fully processed provider configuration from memory. -// This is the primary method called by the account interface and is optimized for minimal latency. +// GetProviderConfigRaw retrieves the raw, unredacted provider configuration from memory. +// This method is for internal use only, particularly by the account implementation. // // Performance characteristics: // - Memory access: ultra-fast direct memory access @@ -393,55 +519,187 @@ func (s *ConfigStore) processMetaConfigEnvVars(rawMetaConfig json.RawMessage, pr // - Thread-safe with read locks for concurrent access // // Returns a copy of the configuration to prevent external modifications. -func (s *ConfigStore) GetProviderConfig(provider schemas.ModelProvider) (*ProviderConfig, error) { +func (s *ConfigStore) GetProviderConfigRaw(provider schemas.ModelProvider) (*ProviderConfig, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + config, exists := s.Providers[provider] + if !exists { + return nil, fmt.Errorf("provider %s not found", provider) + } + + // Return direct reference for maximum performance - this is used by Bifrost core + // CRITICAL: Never modify the returned data as it's shared + return &config, nil +} + +// GetProviderConfigRedacted retrieves a provider configuration with sensitive values redacted. +// This method is intended for external API responses and logging. +// +// The returned configuration has sensitive values redacted: +// - API keys are redacted using RedactKey() +// - Values from environment variables show the original env var name (env.VAR_NAME) +// +// Returns a new copy with redacted values that is safe to expose externally. +func (s *ConfigStore) GetProviderConfigRedacted(provider schemas.ModelProvider) (*ProviderConfig, error) { s.mu.RLock() defer s.mu.RUnlock() - config, exists := s.providers[provider] + config, exists := s.Providers[provider] if !exists { return nil, fmt.Errorf("provider %s not found", provider) } - // Return a copy to prevent external modifications - configCopy := config - return &configCopy, nil + // Create a map for quick lookup of env vars for this provider + envVarsByPath := make(map[string]string) + for envVar, infos := range s.EnvKeys { + for _, info := range infos { + if info.Provider == string(provider) { + envVarsByPath[info.ConfigPath] = envVar + } + } + } + + // Create redacted config with same structure but redacted values + redactedConfig := ProviderConfig{ + NetworkConfig: config.NetworkConfig, + ConcurrencyAndBufferSize: config.ConcurrencyAndBufferSize, + } + + // Create redacted keys + redactedConfig.Keys = make([]schemas.Key, len(config.Keys)) + for i, key := range config.Keys { + redactedConfig.Keys[i] = schemas.Key{ + Models: key.Models, // Copy slice reference - read-only so safe + Weight: key.Weight, + } + + path := fmt.Sprintf("providers.%s.keys[%d]", provider, i) + if envVar, ok := envVarsByPath[path]; ok { + redactedConfig.Keys[i].Value = "env." + envVar + } else { + redactedConfig.Keys[i].Value = RedactKey(key.Value) + } + } + + // Handle meta config redaction if present + if config.MetaConfig != nil { + redactedMetaConfig := s.redactMetaConfig(provider, *config.MetaConfig, envVarsByPath) + redactedConfig.MetaConfig = &redactedMetaConfig + } + + return &redactedConfig, nil +} + +// redactMetaConfig creates a redacted copy of meta config based on provider type +func (s *ConfigStore) redactMetaConfig(provider schemas.ModelProvider, metaConfig schemas.MetaConfig, envVarsByPath map[string]string) schemas.MetaConfig { + switch m := metaConfig.(type) { + case *meta.AzureMetaConfig: + azureConfig := *m // Copy the struct + path := fmt.Sprintf("providers.%s.meta_config.endpoint", provider) + if envVar, ok := envVarsByPath[path]; ok { + azureConfig.Endpoint = "env." + envVar + } else { + azureConfig.Endpoint = RedactKey(azureConfig.Endpoint) + } + if azureConfig.APIVersion != nil { + path = fmt.Sprintf("providers.%s.meta_config.api_version", provider) + if envVar, ok := envVarsByPath[path]; ok { + apiVersion := "env." + envVar + azureConfig.APIVersion = &apiVersion + } + } + return &azureConfig + + case *meta.BedrockMetaConfig: + bedrockConfig := *m // Copy the struct + path := fmt.Sprintf("providers.%s.meta_config.secret_access_key", provider) + if envVar, ok := envVarsByPath[path]; ok { + bedrockConfig.SecretAccessKey = "env." + envVar + } else { + bedrockConfig.SecretAccessKey = RedactKey(bedrockConfig.SecretAccessKey) + } + if bedrockConfig.Region != nil { + path = fmt.Sprintf("providers.%s.meta_config.region", provider) + if envVar, ok := envVarsByPath[path]; ok { + region := "env." + envVar + bedrockConfig.Region = ®ion + } + } + if bedrockConfig.SessionToken != nil { + path = fmt.Sprintf("providers.%s.meta_config.session_token", provider) + if envVar, ok := envVarsByPath[path]; ok { + sessionToken := "env." + envVar + bedrockConfig.SessionToken = &sessionToken + } else { + sessionToken := RedactKey(*bedrockConfig.SessionToken) + bedrockConfig.SessionToken = &sessionToken + } + } + if bedrockConfig.ARN != nil { + path = fmt.Sprintf("providers.%s.meta_config.arn", provider) + if envVar, ok := envVarsByPath[path]; ok { + arn := "env." + envVar + bedrockConfig.ARN = &arn + } + } + return &bedrockConfig + + case *meta.VertexMetaConfig: + vertexConfig := *m // Copy the struct + path := fmt.Sprintf("providers.%s.meta_config.project_id", provider) + if envVar, ok := envVarsByPath[path]; ok { + vertexConfig.ProjectID = "env." + envVar + } + path = fmt.Sprintf("providers.%s.meta_config.region", provider) + if envVar, ok := envVarsByPath[path]; ok { + vertexConfig.Region = "env." + envVar + } + path = fmt.Sprintf("providers.%s.meta_config.auth_credentials", provider) + if envVar, ok := envVarsByPath[path]; ok { + vertexConfig.AuthCredentials = "env." + envVar + } else { + vertexConfig.AuthCredentials = RedactKey(vertexConfig.AuthCredentials) + } + return &vertexConfig + + default: + return metaConfig + } } -// GetAllProviders returns all configured providers. +// GetAllProviders returns all configured provider names. func (s *ConfigStore) GetAllProviders() ([]schemas.ModelProvider, error) { s.mu.RLock() defer s.mu.RUnlock() - providers := make([]schemas.ModelProvider, 0, len(s.providers)) - for provider := range s.providers { + providers := make([]schemas.ModelProvider, 0, len(s.Providers)) + for provider := range s.Providers { providers = append(providers, provider) } return providers, nil } -// UpdateProviderConfig updates a provider configuration in memory with full environment -// variable processing. This method is called when provider configurations are modified -// via the HTTP API and ensures all data processing is done upfront. +// AddProvider adds a new provider configuration to memory with full environment variable +// processing. This method is called when new providers are added via the HTTP API. // // The method: +// - Validates that the provider doesn't already exist // - Processes environment variables in API keys and meta configurations // - Stores the processed configuration in memory // - Updates metadata and timestamps -// - Thread-safe operation with write locks -func (s *ConfigStore) UpdateProviderConfig(provider schemas.ModelProvider, config ProviderConfig) error { +func (s *ConfigStore) AddProvider(provider schemas.ModelProvider, config ProviderConfig) error { s.mu.Lock() defer s.mu.Unlock() - // Process environment variables in keys - for i, key := range config.Keys { - processedValue, err := s.replaceEnvValue(key.Value) - if err != nil { - return fmt.Errorf("failed to process env var in key: %w", err) - } - config.Keys[i].Value = processedValue + // Check if provider already exists + if _, exists := s.Providers[provider]; exists { + return fmt.Errorf("provider %s already exists", provider) } + newEnvKeys := make(map[string]struct{}) + // Process environment variables in meta config if present if config.MetaConfig != nil { rawMetaData, err := json.Marshal(*config.MetaConfig) @@ -449,49 +707,72 @@ func (s *ConfigStore) UpdateProviderConfig(provider schemas.ModelProvider, confi return fmt.Errorf("failed to marshal meta config: %w", err) } - processedMetaData, err := s.processMetaConfigEnvVars(rawMetaData, provider) + processedMetaData, envKeys, err := s.processMetaConfigEnvVars(rawMetaData, provider) + + newEnvKeys = envKeys if err != nil { + s.cleanupEnvKeys(string(provider), "", newEnvKeys) return fmt.Errorf("failed to process env vars in meta config: %w", err) } metaConfig, err := s.parseMetaConfig(processedMetaData, provider) if err != nil { + s.cleanupEnvKeys(string(provider), "", newEnvKeys) return fmt.Errorf("failed to parse processed meta config: %w", err) } config.MetaConfig = metaConfig } - s.providers[provider] = config + // Process environment variables in keys + for i, key := range config.Keys { + processedValue, envVar, err := s.processEnvValue(key.Value) + if err != nil { + s.cleanupEnvKeys(string(provider), "", newEnvKeys) + return fmt.Errorf("failed to process env var in key: %w", err) + } + config.Keys[i].Value = processedValue - s.logger.Info(fmt.Sprintf("Updated configuration for provider: %s", provider)) + // Track environment key if it came from env + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "api_key", + ConfigPath: fmt.Sprintf("providers.%s.keys[%d]", provider, i), + }) + } + } + + s.Providers[provider] = config + + s.logger.Info(fmt.Sprintf("Added provider: %s", provider)) return nil } -// AddProvider adds a new provider configuration to memory with full environment variable -// processing. This method is called when new providers are added via the HTTP API. +// UpdateProviderConfig updates a provider configuration in memory with full environment +// variable processing. This method is called when provider configurations are modified +// via the HTTP API and ensures all data processing is done upfront. // // The method: -// - Validates that the provider doesn't already exist // - Processes environment variables in API keys and meta configurations // - Stores the processed configuration in memory // - Updates metadata and timestamps -func (s *ConfigStore) AddProvider(provider schemas.ModelProvider, config ProviderConfig) error { +// - Thread-safe operation with write locks +// +// Parameters: +// - provider: The provider to update +// - config: The new configuration +// - envKeysToReplace: Map of environment keys that should be replaced (only these will be cleaned up) +func (s *ConfigStore) UpdateProviderConfig(provider schemas.ModelProvider, config ProviderConfig, envKeysToReplace map[string]struct{}) error { s.mu.Lock() defer s.mu.Unlock() - // Check if provider already exists - if _, exists := s.providers[provider]; exists { - return fmt.Errorf("provider %s already exists", provider) - } + // Track new environment variables being added + newEnvKeys := make(map[string]struct{}) - // Process environment variables in keys - for i, key := range config.Keys { - processedValue, err := s.replaceEnvValue(key.Value) - if err != nil { - return fmt.Errorf("failed to process env var in key: %w", err) - } - config.Keys[i].Value = processedValue - } + // Track which old env vars will be replaced (only those specified in envKeysToReplace) + oldEnvKeys := make(map[string]struct{}) // Process environment variables in meta config if present if config.MetaConfig != nil { @@ -500,21 +781,74 @@ func (s *ConfigStore) AddProvider(provider schemas.ModelProvider, config Provide return fmt.Errorf("failed to marshal meta config: %w", err) } - processedMetaData, err := s.processMetaConfigEnvVars(rawMetaData, provider) + // Find old meta config env vars that should be replaced + for envVar, infos := range s.EnvKeys { + for _, info := range infos { + if info.Provider == string(provider) && info.KeyType == "meta_config" { + if _, shouldReplace := envKeysToReplace[envVar]; shouldReplace { + oldEnvKeys[envVar] = struct{}{} + } + } + } + } + + processedMetaData, envKeys, err := s.processMetaConfigEnvVars(rawMetaData, provider) if err != nil { + s.cleanupEnvKeys(string(provider), "", envKeys) // Clean up only new vars on failure return fmt.Errorf("failed to process env vars in meta config: %w", err) } metaConfig, err := s.parseMetaConfig(processedMetaData, provider) if err != nil { + s.cleanupEnvKeys(string(provider), "", envKeys) // Clean up only new vars on failure return fmt.Errorf("failed to parse processed meta config: %w", err) } config.MetaConfig = metaConfig + + // Add the new env vars to tracking + for envVar := range envKeys { + newEnvKeys[envVar] = struct{}{} + } } - s.providers[provider] = config + // Find old API key env vars that should be replaced + for envVar, infos := range s.EnvKeys { + for _, info := range infos { + if info.Provider == string(provider) && info.KeyType == "api_key" { + if _, shouldReplace := envKeysToReplace[envVar]; shouldReplace { + oldEnvKeys[envVar] = struct{}{} + } + } + } + } - s.logger.Info(fmt.Sprintf("Added provider: %s", provider)) + // Process environment variables in keys + for i, key := range config.Keys { + processedValue, envVar, err := s.processEnvValue(key.Value) + if err != nil { + s.cleanupEnvKeys(string(provider), "", newEnvKeys) // Clean up only new vars on failure + return fmt.Errorf("failed to process env var in key: %w", err) + } + config.Keys[i].Value = processedValue + + // Track environment key if it came from env + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: string(provider), + KeyType: "api_key", + ConfigPath: fmt.Sprintf("providers.%s.keys[%d]", provider, i), + }) + } + } + + s.Providers[provider] = config + + // Clean up old env vars that were replaced + s.cleanupEnvKeys(string(provider), "", oldEnvKeys) + + s.logger.Info(fmt.Sprintf("Updated configuration for provider: %s", provider)) return nil } @@ -523,11 +857,12 @@ func (s *ConfigStore) RemoveProvider(provider schemas.ModelProvider) error { s.mu.Lock() defer s.mu.Unlock() - if _, exists := s.providers[provider]; !exists { + if _, exists := s.Providers[provider]; !exists { return fmt.Errorf("provider %s not found", provider) } - delete(s.providers, provider) + delete(s.Providers, provider) + s.cleanupEnvKeys(string(provider), "", nil) s.logger.Info(fmt.Sprintf("Removed provider: %s", provider)) return nil @@ -542,19 +877,36 @@ func (s *ConfigStore) RemoveProvider(provider schemas.ModelProvider) error { // // Returns an error if any required environment variable is missing. // This approach ensures type safety while supporting environment variable substitution. -func (s *ConfigStore) processMCPEnvVars() { +func (s *ConfigStore) processMCPEnvVars() error { + var missingEnvVars []string + // Process each client config - for i, clientConfig := range s.mcpConfig.ClientConfigs { + for i, clientConfig := range s.MCPConfig.ClientConfigs { // Process ConnectionString if present if clientConfig.ConnectionString != nil { - newValue, err := s.replaceEnvValue(*clientConfig.ConnectionString) + newValue, envVar, err := s.processEnvValue(*clientConfig.ConnectionString) if err != nil { s.logger.Warn(fmt.Sprintf("failed to process env vars in MCP client %s: %v", clientConfig.Name, err)) + missingEnvVars = append(missingEnvVars, envVar) continue } - s.mcpConfig.ClientConfigs[i].ConnectionString = &newValue + if envVar != "" { + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: "", + KeyType: "connection_string", + ConfigPath: fmt.Sprintf("mcp.client_configs[%d].connection_string", i), + }) + } + s.MCPConfig.ClientConfigs[i].ConnectionString = &newValue } } + + if len(missingEnvVars) > 0 { + return fmt.Errorf("missing environment variables: %v", missingEnvVars) + } + + return nil } // SetBifrostClient sets the Bifrost client in the store. @@ -567,13 +919,6 @@ func (s *ConfigStore) SetBifrostClient(client *bifrost.Bifrost) { s.client = client } -// GetMCPConfig retrieves the processed MCP configuration from memory. -// Returns nil if no MCP configuration was loaded. -// The returned configuration has all environment variables already processed. -func (s *ConfigStore) GetMCPConfig() *schemas.MCPConfig { - return s.mcpConfig -} - // AddMCPClient adds a new MCP client to the configuration. // This method is called when a new MCP client is added via the HTTP API. // @@ -589,16 +934,38 @@ func (s *ConfigStore) AddMCPClient(clientConfig schemas.MCPClientConfig) error { s.muMCP.Lock() defer s.muMCP.Unlock() - if s.mcpConfig == nil { - s.mcpConfig = &schemas.MCPConfig{} + if s.MCPConfig == nil { + s.MCPConfig = &schemas.MCPConfig{} } - s.mcpConfig.ClientConfigs = append(s.mcpConfig.ClientConfigs, clientConfig) - s.processMCPEnvVars() + // Track new environment variables + newEnvKeys := make(map[string]struct{}) + + s.MCPConfig.ClientConfigs = append(s.MCPConfig.ClientConfigs, clientConfig) + + // Process environment variables in the new client config + if clientConfig.ConnectionString != nil { + processedValue, envVar, err := s.processEnvValue(*clientConfig.ConnectionString) + if err != nil { + s.MCPConfig.ClientConfigs = s.MCPConfig.ClientConfigs[:len(s.MCPConfig.ClientConfigs)-1] + return fmt.Errorf("failed to process env var in connection string: %w", err) + } + if envVar != "" { + newEnvKeys[envVar] = struct{}{} + s.EnvKeys[envVar] = append(s.EnvKeys[envVar], EnvKeyInfo{ + EnvVar: envVar, + Provider: "", + KeyType: "connection_string", + ConfigPath: fmt.Sprintf("mcp.client_configs.%s.connection_string", clientConfig.Name), + }) + } + s.MCPConfig.ClientConfigs[len(s.MCPConfig.ClientConfigs)-1].ConnectionString = &processedValue + } // Config with processed env vars - if err := s.client.AddMCPClient(s.mcpConfig.ClientConfigs[len(s.mcpConfig.ClientConfigs)-1]); err != nil { - s.mcpConfig.ClientConfigs = s.mcpConfig.ClientConfigs[:len(s.mcpConfig.ClientConfigs)-1] + if err := s.client.AddMCPClient(s.MCPConfig.ClientConfigs[len(s.MCPConfig.ClientConfigs)-1]); err != nil { + s.MCPConfig.ClientConfigs = s.MCPConfig.ClientConfigs[:len(s.MCPConfig.ClientConfigs)-1] + s.cleanupEnvKeys("", clientConfig.Name, newEnvKeys) return fmt.Errorf("failed to add MCP client: %w", err) } @@ -620,7 +987,7 @@ func (s *ConfigStore) RemoveMCPClient(name string) error { s.muMCP.Lock() defer s.muMCP.Unlock() - if s.mcpConfig == nil { + if s.MCPConfig == nil { return fmt.Errorf("no MCP config found") } @@ -628,13 +995,15 @@ func (s *ConfigStore) RemoveMCPClient(name string) error { return fmt.Errorf("failed to remove MCP client: %w", err) } - for i, clientConfig := range s.mcpConfig.ClientConfigs { + for i, clientConfig := range s.MCPConfig.ClientConfigs { if clientConfig.Name == name { - s.mcpConfig.ClientConfigs = append(s.mcpConfig.ClientConfigs[:i], s.mcpConfig.ClientConfigs[i+1:]...) + s.MCPConfig.ClientConfigs = append(s.MCPConfig.ClientConfigs[:i], s.MCPConfig.ClientConfigs[i+1:]...) break } } + s.cleanupEnvKeys("", name, nil) + return nil } @@ -653,7 +1022,7 @@ func (s *ConfigStore) EditMCPClientTools(name string, toolsToAdd []string, tools s.muMCP.Lock() defer s.muMCP.Unlock() - if s.mcpConfig == nil { + if s.MCPConfig == nil { return fmt.Errorf("no MCP config found") } @@ -661,13 +1030,144 @@ func (s *ConfigStore) EditMCPClientTools(name string, toolsToAdd []string, tools return fmt.Errorf("failed to edit MCP client tools: %w", err) } - for i, clientConfig := range s.mcpConfig.ClientConfigs { + for i, clientConfig := range s.MCPConfig.ClientConfigs { if clientConfig.Name == name { - s.mcpConfig.ClientConfigs[i].ToolsToExecute = toolsToAdd - s.mcpConfig.ClientConfigs[i].ToolsToSkip = toolsToRemove + s.MCPConfig.ClientConfigs[i].ToolsToExecute = toolsToAdd + s.MCPConfig.ClientConfigs[i].ToolsToSkip = toolsToRemove break } } return nil } + +// RedactMCPClientConfig creates a redacted copy of an MCP client configuration. +// Connection strings are either redacted or replaced with their environment variable names. +func (s *ConfigStore) RedactMCPClientConfig(config schemas.MCPClientConfig) schemas.MCPClientConfig { + // Create a copy with basic fields + configCopy := schemas.MCPClientConfig{ + Name: config.Name, + ConnectionType: config.ConnectionType, + ConnectionString: config.ConnectionString, + StdioConfig: config.StdioConfig, + ToolsToExecute: append([]string{}, config.ToolsToExecute...), + ToolsToSkip: append([]string{}, config.ToolsToSkip...), + } + + // Handle connection string if present + if config.ConnectionString != nil { + connStr := *config.ConnectionString + + // Check if this value came from an env var + for envVar, infos := range s.EnvKeys { + for _, info := range infos { + if info.Provider == "" && info.KeyType == "connection_string" && info.ConfigPath == fmt.Sprintf("mcp.client_configs.%s.connection_string", config.Name) { + connStr = "env." + envVar + break + } + } + } + + // If not from env var, redact it + if !strings.HasPrefix(connStr, "env.") { + connStr = RedactKey(connStr) + } + configCopy.ConnectionString = &connStr + } + + return configCopy +} + +// RedactKey redacts sensitive key values by showing only the first and last 4 characters +func RedactKey(key string) string { + if key == "" { + return "" + } + + // If key is 8 characters or less, just return all asterisks + if len(key) <= 8 { + return strings.Repeat("*", len(key)) + } + + // Show first 4 and last 4 characters, replace middle with asterisks + prefix := key[:4] + suffix := key[len(key)-4:] + middle := strings.Repeat("*", 24) + + return prefix + middle + suffix +} + +// IsRedacted checks if a key value is redacted, either by being an environment variable +// reference (env.VAR_NAME) or containing the exact redaction pattern from RedactKey. +func IsRedacted(key string) bool { + if key == "" { + return false + } + + // Check if it's an environment variable reference + if strings.HasPrefix(key, "env.") { + return true + } + + // Check for exact redaction pattern: 4 chars + 24 asterisks + 4 chars + if len(key) == 32 { + middle := key[4:28] + if middle == strings.Repeat("*", 24) { + return true + } + } + + return false +} + +// cleanupEnvKeys removes environment variable entries from the store based on the given criteria. +// If envVarsToRemove is nil, it removes all env vars for the specified provider/client. +// If envVarsToRemove is provided, it only removes those specific env vars. +// +// Parameters: +// - provider: Provider name to clean up (empty string for MCP clients) +// - mcpClientName: MCP client name to clean up (empty string for providers) +// - envVarsToRemove: Optional map of specific env vars to remove (nil to remove all) +func (s *ConfigStore) cleanupEnvKeys(provider string, mcpClientName string, envVarsToRemove map[string]struct{}) { + // If envVarsToRemove is provided, only clean those specific vars + if envVarsToRemove != nil { + for envVar := range envVarsToRemove { + s.cleanupEnvVar(envVar, provider, mcpClientName) + } + return + } + + // If envVarsToRemove is nil, clean all vars for the provider/client + for envVar := range s.EnvKeys { + s.cleanupEnvVar(envVar, provider, mcpClientName) + } +} + +// cleanupEnvVar removes entries for a specific environment variable based on provider/client. +// This is a helper function to avoid duplicating the filtering logic. +func (s *ConfigStore) cleanupEnvVar(envVar, provider, mcpClientName string) { + infos := s.EnvKeys[envVar] + if len(infos) == 0 { + return + } + + // Keep entries that don't match the provider/client we're cleaning up + filteredInfos := make([]EnvKeyInfo, 0, len(infos)) + for _, info := range infos { + shouldKeep := false + if provider != "" { + shouldKeep = info.Provider != provider + } else if mcpClientName != "" { + shouldKeep = info.Provider != "" || !strings.HasPrefix(info.ConfigPath, fmt.Sprintf("mcp.client_configs.%s", mcpClientName)) + } + if shouldKeep { + filteredInfos = append(filteredInfos, info) + } + } + + if len(filteredInfos) == 0 { + delete(s.EnvKeys, envVar) + } else { + s.EnvKeys[envVar] = filteredInfos + } +} diff --git a/transports/bifrost-http/main.go b/transports/bifrost-http/main.go index a2b37228f2..5157247668 100644 --- a/transports/bifrost-http/main.go +++ b/transports/bifrost-http/main.go @@ -50,7 +50,6 @@ package main import ( - "encoding/json" "flag" "fmt" "log" @@ -74,10 +73,9 @@ import ( // Command line flags var ( - initialPoolSize int // Initial size of the connection pool - port string // Port to run the server on - configPath string // Path to the config file - pluginsToLoad []string // Path to the plugins + port string // Port to run the server on + configPath string // Path to the config file + pluginsToLoad []string // Path to the plugins ) // init initializes command line flags and validates required configuration. @@ -88,7 +86,6 @@ var ( func init() { pluginString := "" - flag.IntVar(&initialPoolSize, "pool-size", 300, "Initial pool size for Bifrost") flag.StringVar(&port, "port", "8080", "Port to run the server on") flag.StringVar(&configPath, "config", "", "Path to the config file") flag.StringVar(&pluginString, "plugins", "", "Comma separated list of plugins to load") @@ -131,28 +128,6 @@ func main() { logger := bifrost.NewDefaultLogger(schemas.LogLevelInfo) - // Define a struct to unmarshal the entire config file - var config struct { - Client struct { - DropExcessRequests bool `json:"drop_excess_requests"` - PrometheusLabels []string `json:"prometheus_labels"` - } `json:"client"` - Providers json.RawMessage `json:"providers"` - MCP *schemas.MCPConfig `json:"mcp"` - } - - // Read and parse config - data, err := os.ReadFile(configPath) - if err != nil { - log.Fatalf("failed to read config file: %v", err) - } - if err := json.Unmarshal(data, &config); err != nil { - log.Fatalf("failed to parse config JSON: %v", err) - } - - telemetry.InitPrometheusMetrics(config.Client.PrometheusLabels) - log.Println("Prometheus Go/Process collectors registered.") - // Initialize high-performance configuration store with caching store, err := lib.NewConfigStore(logger) if err != nil { @@ -169,10 +144,6 @@ func main() { // The account interface now benefits from ultra-fast config access times via in-memory storage account := lib.NewBaseAccount(store) - // Get the processed MCP configuration from the store - // All environment variable processing is already done during LoadFromConfig - mcpConfig := store.GetMCPConfig() - loadedPlugins := []schemas.Plugin{} for _, plugin := range pluginsToLoad { @@ -197,6 +168,9 @@ func main() { } } + telemetry.InitPrometheusMetrics(store.ClientConfig.PrometheusLabels) + log.Println("Prometheus Go/Process collectors registered.") + promPlugin := telemetry.NewPrometheusPlugin() loggingPlugin, err := logging.NewLoggerPlugin(nil) if err != nil { @@ -207,10 +181,10 @@ func main() { client, err := bifrost.Init(schemas.BifrostConfig{ Account: account, - InitialPoolSize: initialPoolSize, - DropExcessRequests: config.Client.DropExcessRequests, + InitialPoolSize: store.ClientConfig.InitialPoolSize, + DropExcessRequests: store.ClientConfig.DropExcessRequests, Plugins: loadedPlugins, - MCPConfig: mcpConfig, + MCPConfig: store.MCPConfig, Logger: logger, }) if err != nil { @@ -224,7 +198,7 @@ func main() { completionHandler := handlers.NewCompletionHandler(client, logger) mcpHandler := handlers.NewMCPHandler(client, logger, store) integrationHandler := handlers.NewIntegrationHandler(client) - configHandler := handlers.NewConfigHandler(client, logger, configPath) + configHandler := handlers.NewConfigHandler(client, logger, store, configPath) loggingHandler := handlers.NewLoggingHandler(loggingPlugin.GetPluginLogManager(), logger) wsHandler := handlers.NewWebSocketHandler(loggingPlugin.GetPluginLogManager(), logger)