From d83c822c7354b64fb34073dcc85e7d079c74b9e9 Mon Sep 17 00:00:00 2001
From: EnumDev <enumdev@enumerated.dev>
Date: Wed, 24 Sep 2025 13:34:46 +0300
Subject: [PATCH] Minor enabled service stage fixes

---
 src/esvm/service.go | 13 ++++++++-----
 src/esvm/socket.go  |  7 ++++++-
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/esvm/service.go b/src/esvm/service.go
index 82b3b27..7875d5b 100644
--- a/src/esvm/service.go
+++ b/src/esvm/service.go
@@ -346,14 +346,17 @@ func (service *EnitService) SetEnabled(stage int) error {
 	}
 
 	// Remove service from current stage
-	if s != 0 {
-		EnabledServices[s] = slices.DeleteFunc(EnabledServices[s], func(name string) bool {
-			return name == service.Name
-		})
+	EnabledServices[s] = slices.DeleteFunc(EnabledServices[s], func(name string) bool {
+		return name == service.Name
+	})
+	if len(EnabledServices[s]) == 0 {
+		delete(EnabledServices, s)
 	}
 
 	// Add service to stage
-	EnabledServices[stage] = append(EnabledServices[stage], service.Name)
+	if stage != 0 {
+		EnabledServices[stage] = append(EnabledServices[stage], service.Name)
+	}
 
 	// Save enabled services to file
 	data, err := yaml.Marshal(EnabledServices)
diff --git a/src/esvm/socket.go b/src/esvm/socket.go
index 588fe36..a9854ed 100644
--- a/src/esvm/socket.go
+++ b/src/esvm/socket.go
@@ -197,7 +197,12 @@ func handleSetEnabledServiceCommand(conn net.Conn, jsonData map[string]any) {
 		return
 	}
 
-	conn.Write(wrapSuccessMsgInJson(fmt.Sprintf("Service (%s) was enabled sucessfully", serviceName.(string))))
+	if serviceStage == 0 {
+		conn.Write(wrapSuccessMsgInJson(fmt.Sprintf("Service (%s) was disabled sucessfully", serviceName.(string))))
+	} else {
+		conn.Write(wrapSuccessMsgInJson(fmt.Sprintf("Service (%s) was enabled sucessfully", serviceName.(string))))
+	}
+
 }
 
 func handleStatusServiceCommand(conn net.Conn, jsonData map[string]any) {