diff --git a/pkg/config/config.go b/pkg/config/config.go index 77637b8..6b34b63 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -41,7 +41,7 @@ type AppConfig struct { Backends BackendConfig `yaml:"backends"` Instances InstancesConfig `yaml:"instances"` Auth AuthConfig `yaml:"auth"` - Nodes []NodeConfig `yaml:"nodes,omitempty"` + Nodes map[string]NodeConfig `yaml:"nodes,omitempty"` Version string `yaml:"-"` CommitHash string `yaml:"-"` BuildTime string `yaml:"-"` @@ -130,7 +130,6 @@ type AuthConfig struct { } type NodeConfig struct { - Name string `yaml:"name"` Address string `yaml:"address"` APIKey string `yaml:"api_key,omitempty"` } diff --git a/pkg/manager/manager.go b/pkg/manager/manager.go index a5a7138..37989db 100644 --- a/pkg/manager/manager.go +++ b/pkg/manager/manager.go @@ -63,15 +63,16 @@ type instanceManager struct { } // NewInstanceManager creates a new instance of InstanceManager. -func NewInstanceManager(backendsConfig config.BackendConfig, instancesConfig config.InstancesConfig, nodesConfig []config.NodeConfig) InstanceManager { +func NewInstanceManager(backendsConfig config.BackendConfig, instancesConfig config.InstancesConfig, nodesConfig map[string]config.NodeConfig) InstanceManager { if instancesConfig.TimeoutCheckInterval <= 0 { instancesConfig.TimeoutCheckInterval = 5 // Default to 5 minutes if not set } // Build node config map for quick lookup nodeConfigMap := make(map[string]*config.NodeConfig) - for i := range nodesConfig { - nodeConfigMap[nodesConfig[i].Name] = &nodesConfig[i] + for name := range nodesConfig { + nodeCopy := nodesConfig[name] + nodeConfigMap[name] = &nodeCopy } im := &instanceManager{ diff --git a/pkg/manager/manager_test.go b/pkg/manager/manager_test.go index 5110e52..e59e2eb 100644 --- a/pkg/manager/manager_test.go +++ b/pkg/manager/manager_test.go @@ -34,7 +34,7 @@ func TestNewInstanceManager(t *testing.T) { TimeoutCheckInterval: 5, } - mgr := manager.NewInstanceManager(backendConfig, cfg, []config.NodeConfig{}) + mgr := manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) if mgr == nil { t.Fatal("NewInstanceManager returned nil") } @@ -69,7 +69,7 @@ func TestPersistence(t *testing.T) { } // Test instance persistence on creation - manager1 := manager.NewInstanceManager(backendConfig, cfg, []config.NodeConfig{}) + manager1 := manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) options := &instance.CreateInstanceOptions{ BackendType: backends.BackendTypeLlamaCpp, LlamaServerOptions: &llamacpp.LlamaServerOptions{ @@ -90,7 +90,7 @@ func TestPersistence(t *testing.T) { } // Test loading instances from disk - manager2 := manager.NewInstanceManager(backendConfig, cfg, []config.NodeConfig{}) + manager2 := manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) instances, err := manager2.ListInstances() if err != nil { t.Fatalf("ListInstances failed: %v", err) @@ -207,7 +207,7 @@ func createTestManager() manager.InstanceManager { DefaultRestartDelay: 5, TimeoutCheckInterval: 5, } - return manager.NewInstanceManager(backendConfig, cfg, []config.NodeConfig{}) + return manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) } func TestAutoRestartDisabledInstanceStatus(t *testing.T) { @@ -227,7 +227,7 @@ func TestAutoRestartDisabledInstanceStatus(t *testing.T) { } // Create first manager and instance with auto-restart disabled - manager1 := manager.NewInstanceManager(backendConfig, cfg, []config.NodeConfig{}) + manager1 := manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) autoRestart := false options := &instance.CreateInstanceOptions{ @@ -252,7 +252,7 @@ func TestAutoRestartDisabledInstanceStatus(t *testing.T) { manager1.Shutdown() // Create second manager (simulating restart of llamactl) - manager2 := manager.NewInstanceManager(backendConfig, cfg, []config.NodeConfig{}) + manager2 := manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) // Get the loaded instance loadedInst, err := manager2.GetInstance("test-instance") diff --git a/pkg/manager/operations_test.go b/pkg/manager/operations_test.go index da26742..fdeb44f 100644 --- a/pkg/manager/operations_test.go +++ b/pkg/manager/operations_test.go @@ -75,7 +75,7 @@ func TestCreateInstance_ValidationAndLimits(t *testing.T) { MaxInstances: 1, // Very low limit for testing TimeoutCheckInterval: 5, } - limitedManager := manager.NewInstanceManager(backendConfig, cfg, nil) + limitedManager := manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) _, err = limitedManager.CreateInstance("instance1", options) if err != nil { diff --git a/pkg/manager/timeout_test.go b/pkg/manager/timeout_test.go index 31b4298..55cd781 100644 --- a/pkg/manager/timeout_test.go +++ b/pkg/manager/timeout_test.go @@ -23,7 +23,7 @@ func TestTimeoutFunctionality(t *testing.T) { MaxInstances: 5, } - manager := manager.NewInstanceManager(backendConfig, cfg, nil) + manager := manager.NewInstanceManager(backendConfig, cfg, map[string]config.NodeConfig{}) if manager == nil { t.Fatal("Manager should be initialized with timeout checker") } diff --git a/pkg/server/handlers_instances.go b/pkg/server/handlers_instances.go index 8c325ae..88e974f 100644 --- a/pkg/server/handlers_instances.go +++ b/pkg/server/handlers_instances.go @@ -3,7 +3,6 @@ package server import ( "encoding/json" "fmt" - "llamactl/pkg/config" "llamactl/pkg/instance" "llamactl/pkg/manager" "net/http" @@ -414,15 +413,8 @@ func (h *Handler) RemoteInstanceProxy(w http.ResponseWriter, r *http.Request, na if !exists { // Find node configuration - var nodeConfig *config.NodeConfig - for i := range h.cfg.Nodes { - if h.cfg.Nodes[i].Name == nodeName { - nodeConfig = &h.cfg.Nodes[i] - break - } - } - - if nodeConfig == nil { + nodeConfig, exists := h.cfg.Nodes[nodeName] + if !exists { http.Error(w, fmt.Sprintf("Node %s not found", nodeName), http.StatusInternalServerError) return } diff --git a/pkg/server/handlers_nodes.go b/pkg/server/handlers_nodes.go index 7d1116f..98a4b43 100644 --- a/pkg/server/handlers_nodes.go +++ b/pkg/server/handlers_nodes.go @@ -2,7 +2,6 @@ package server import ( "encoding/json" - "llamactl/pkg/config" "net/http" "github.com/go-chi/chi/v5" @@ -10,26 +9,24 @@ import ( // NodeResponse represents a sanitized node configuration for API responses type NodeResponse struct { - Name string `json:"name"` Address string `json:"address"` } // ListNodes godoc // @Summary List all configured nodes -// @Description Returns a list of all nodes configured in the server +// @Description Returns a map of all nodes configured in the server (node name -> node config) // @Tags nodes // @Security ApiKeyAuth // @Produces json -// @Success 200 {array} NodeResponse "List of nodes" +// @Success 200 {object} map[string]NodeResponse "Map of nodes" // @Failure 500 {string} string "Internal Server Error" // @Router /nodes [get] func (h *Handler) ListNodes() http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { - // Convert to sanitized response format - nodeResponses := make([]NodeResponse, len(h.cfg.Nodes)) - for i, node := range h.cfg.Nodes { - nodeResponses[i] = NodeResponse{ - Name: node.Name, + // Convert to sanitized response format (map of name -> NodeResponse) + nodeResponses := make(map[string]NodeResponse, len(h.cfg.Nodes)) + for name, node := range h.cfg.Nodes { + nodeResponses[name] = NodeResponse{ Address: node.Address, } } @@ -62,22 +59,14 @@ func (h *Handler) GetNode() http.HandlerFunc { return } - var nodeConfig *config.NodeConfig - for i := range h.cfg.Nodes { - if h.cfg.Nodes[i].Name == name { - nodeConfig = &h.cfg.Nodes[i] - break - } - } - - if nodeConfig == nil { + nodeConfig, exists := h.cfg.Nodes[name] + if !exists { http.Error(w, "Node not found", http.StatusNotFound) return } // Convert to sanitized response format nodeResponse := NodeResponse{ - Name: nodeConfig.Name, Address: nodeConfig.Address, } diff --git a/pkg/server/handlers_openai.go b/pkg/server/handlers_openai.go index eea8440..c6e56e9 100644 --- a/pkg/server/handlers_openai.go +++ b/pkg/server/handlers_openai.go @@ -5,7 +5,6 @@ import ( "encoding/json" "fmt" "io" - "llamactl/pkg/config" "llamactl/pkg/instance" "net/http" "net/http/httputil" @@ -170,15 +169,8 @@ func (h *Handler) RemoteOpenAIProxy(w http.ResponseWriter, r *http.Request, mode if !exists { // Find node configuration - var nodeConfig *config.NodeConfig - for i := range h.cfg.Nodes { - if h.cfg.Nodes[i].Name == nodeName { - nodeConfig = &h.cfg.Nodes[i] - break - } - } - - if nodeConfig == nil { + nodeConfig, exists := h.cfg.Nodes[nodeName] + if !exists { http.Error(w, fmt.Sprintf("Node %s not found", nodeName), http.StatusInternalServerError) return }