From da1d88a7bd8669eee7682d0ca957603691ef11a4 Mon Sep 17 00:00:00 2001
From: yuyufei <yuyufei@archebase.ai>
Date: Sat, 25 Jul 2026 14:44:06 +0800
Subject: [PATCH] fix: serve alb health on routed prefixes

---
 deploy/helm/keystone-stack/README.md |  6 ++---
 internal/server/server.go            | 18 +++++++++++++
 internal/server/server_test.go       | 38 ++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/deploy/helm/keystone-stack/README.md b/deploy/helm/keystone-stack/README.md
index 09f7cc5..5987db3 100644
--- a/deploy/helm/keystone-stack/README.md
+++ b/deploy/helm/keystone-stack/README.md
@@ -169,9 +169,9 @@ Synapse uses the same-origin `/api/v1` path. The chart's HTTP Ingress sends
 `/api` and `/swagger` to Keystone, `/transfer` and `/recorder` to Keystone's
 WebSocket ports, and all remaining paths to Synapse.
 
-Keystone also serves a root health response for ALB backend health checks. The
-public root path still routes to Synapse because the HTTP Ingress sends `/` to
-the Synapse Service.
+Keystone also serves root and `/api` health responses for ALB backend health
+checks. The public root path still routes to Synapse because the HTTP Ingress
+sends `/` to the Synapse Service.
 
 The gRPC Ingress is separate so the VKE ALB annotations can select listener
 `50053` and backend protocol `grpc` without affecting browser/API routing.
diff --git a/internal/server/server.go b/internal/server/server.go
index 0db3543..a8dd33a 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -75,6 +75,20 @@ func axonTransferWriteTimeout(cfg *config.TransferConfig) time.Duration {
 	return time.Duration(cfg.WriteTimeout) * time.Second
 }
 
+func loadBalancerHealthHandler(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodGet && r.Method != http.MethodHead {
+		w.Header().Set("Allow", "GET, HEAD")
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		return
+	}
+
+	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+	w.WriteHeader(http.StatusOK)
+	if r.Method != http.MethodHead {
+		_, _ = w.Write([]byte("ok\n"))
+	}
+}
+
 // New creates a new server instance.
 // db and s3Client are optional; pass nil to disable Verified ACK.
 // syncWorker is optional; pass nil to disable cloud sync APIs.
@@ -231,6 +245,8 @@ func (s *Server) buildRoutes() http.Handler {
 
 	// Root health check for load balancers that probe "/" on the Keystone backend.
 	s.engine.GET("/", s.health.Handler)
+	// Prefix health check for load balancers that probe the Ingress backend path.
+	s.engine.GET("/api", s.health.Handler)
 
 	// Health check
 	s.health.RegisterAPI(v1)
@@ -344,6 +360,7 @@ func (s *Server) buildRoutes() http.Handler {
 func (s *Server) buildTransferWSRoutes(transferHandler *handlers.TransferHandler) http.Handler {
 	mux := http.NewServeMux()
 
+	mux.HandleFunc("/transfer", loadBalancerHealthHandler)
 	mux.HandleFunc("/transfer/", func(w http.ResponseWriter, r *http.Request) {
 		// Extract device_id from URL path
 		deviceID := strings.TrimPrefix(r.URL.Path, "/transfer/")
@@ -361,6 +378,7 @@ func (s *Server) buildTransferWSRoutes(transferHandler *handlers.TransferHandler
 func (s *Server) buildRecorderWSRoutes(recorderHandler *handlers.RecorderHandler) http.Handler {
 	mux := http.NewServeMux()
 
+	mux.HandleFunc("/recorder", loadBalancerHealthHandler)
 	mux.HandleFunc("/recorder/", func(w http.ResponseWriter, r *http.Request) {
 		deviceID := strings.TrimPrefix(r.URL.Path, "/recorder/")
 		if deviceID == "" || deviceID == r.URL.Path {
diff --git a/internal/server/server_test.go b/internal/server/server_test.go
index fc77910..fad94e3 100644
--- a/internal/server/server_test.go
+++ b/internal/server/server_test.go
@@ -51,6 +51,7 @@ func TestHTTPHealthRoutes(t *testing.T) {
 
 	tests := []string{
 		"/",
+		"/api",
 		"/api/v1/health",
 	}
 	for _, path := range tests {
@@ -66,3 +67,40 @@ func TestHTTPHealthRoutes(t *testing.T) {
 		})
 	}
 }
+
+func TestWebSocketHealthRoutes(t *testing.T) {
+	srv := New(&config.Config{
+		Server: config.ServerConfig{
+			BindAddr: ":8080",
+		},
+		AxonTransfer: config.TransferConfig{
+			WSPort:    8090,
+			MaxEvents: 10,
+		},
+		AxonRecorder: config.RecorderConfig{
+			WSPort:          8091,
+			ResponseTimeout: 1,
+		},
+	}, nil, nil, nil)
+
+	tests := []struct {
+		name    string
+		handler http.Handler
+		path    string
+	}{
+		{name: "transfer", handler: srv.transferWSServer.Handler, path: "/transfer"},
+		{name: "recorder", handler: srv.recorderWSServer.Handler, path: "/recorder"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, tt.path, nil)
+			w := httptest.NewRecorder()
+
+			tt.handler.ServeHTTP(w, req)
+
+			if w.Code != http.StatusOK {
+				t.Fatalf("GET %s status=%d want=%d body=%s", tt.path, w.Code, http.StatusOK, w.Body.String())
+			}
+		})
+	}
+}