Add health command for hub and align agent health command

beszel/cmd/agent/agent.go

@@ -6,7 +6,6 @@ import (
 	"flag"
 	"fmt"
 	"log"
-	"log/slog"
 	"os"
 
 	"golang.org/x/crypto/ssh"
@@ -56,9 +55,12 @@ func (opts *cmdOptions) parse() bool {
 		flag.CommandLine.Parse(args)
 		addr := opts.getAddress()
 		network := agent.GetNetwork(addr)
-		exitCode, err := agent.Health(addr, network)
-		slog.Info("Health", "code", exitCode, "err", err)
-		os.Exit(exitCode)
+		err := agent.Health(addr, network)
+		if err != nil {
+			log.Fatal(err)
+		}
+		fmt.Print("ok")
+		return true
 	}
 
 	flag.Parse()

beszel/cmd/hub/hub.go

@@ -4,7 +4,10 @@ import (
 	"beszel"
 	"beszel/internal/hub"
 	_ "beszel/migrations"
+	"fmt"
+	"net/http"
 	"os"
+	"time"
 
 	"github.com/pocketbase/pocketbase"
 	"github.com/pocketbase/pocketbase/plugins/migratecmd"
@@ -12,6 +15,16 @@ import (
 )
 
 func main() {
+	// handle health check first to prevent unneeded execution
+	if len(os.Args) > 3 && os.Args[1] == "health" {
+		url := os.Args[3]
+		if err := checkHealth(url); err != nil {
+			log.Fatal(err)
+		}
+		fmt.Print("ok")
+		return
+	}
+
 	baseApp := getBaseApp()
 	h := hub.NewHub(baseApp)
 	h.BootstrapHub()
@@ -35,6 +48,8 @@ func getBaseApp() *pocketbase.PocketBase {
 		Short: "Update " + beszel.AppName + " to the latest version",
 		Run:   hub.Update,
 	})
+	// add health command
+	baseApp.RootCmd.AddCommand(newHealthCmd())
 
 	// enable auto creation of migration files when making collection changes in the Admin UI
 	migratecmd.MustRegister(baseApp, baseApp.RootCmd, migratecmd.Config{
@@ -44,3 +59,39 @@ func getBaseApp() *pocketbase.PocketBase {
 
 	return baseApp
 }
+
+func newHealthCmd() *cobra.Command {
+	var baseURL string
+
+	healthCmd := &cobra.Command{
+		Use:   "health",
+		Short: "Check health of running hub",
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := checkHealth(baseURL); err != nil {
+				log.Fatal(err)
+			}
+			os.Exit(0)
+		},
+	}
+	healthCmd.Flags().StringVar(&baseURL, "url", "", "base URL")
+	healthCmd.MarkFlagRequired("url")
+	return healthCmd
+}
+
+// checkHealth checks the health of the hub.
+func checkHealth(baseURL string) error {
+	client := &http.Client{
+		Timeout: time.Second * 3,
+	}
+	healthURL := baseURL + "/api/health"
+	resp, err := client.Get(healthURL)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		return fmt.Errorf("%s returned status %d", healthURL, resp.StatusCode)
+	}
+	return nil
+}