feat: wire manifest readinessProbe to remote session controller

api/v1alpha1/amaltheasession_children.go

@@ -282,6 +282,10 @@ func (cr *AmaltheaSession) Service() v1.Service {
 		},
 	}
 	if cr.Spec.SessionLocation == Remote {
+		// NOTE: In order to connect through the reverse tunnel,
+		// we need to have the tunnel established so we publish
+		// the service without the pod being ready
+		svc.Spec.PublishNotReadyAddresses = true
 		svc.Spec.Ports = append(svc.Spec.Ports, v1.ServicePort{
 			Protocol:   v1.ProtocolTCP,
 			Name:       tunnelServiceName,
@@ -991,6 +995,18 @@ func (cr *AmaltheaSession) sessionContainerRemote(volumeMounts []v1.VolumeMount)
 			Name:  "RSC_SERVER_PORT",
 			Value: fmt.Sprintf("%d", RemoteSessionControllerPort),
 		},
+		v1.EnvVar{
+			Name:  "RSC_SESSION_PORT",
+			Value: fmt.Sprintf("%d", cr.Spec.Session.Port),
+		},
+		v1.EnvVar{
+			Name:  "RSC_SESSION_URL_PATH",
+			Value: cr.Spec.Session.URLPath,
+		},
+		v1.EnvVar{
+			Name:  "RSC_READINESS_PROBE_TYPE",
+			Value: string(cr.Spec.Session.ReadinessProbe.Type),
+		},
 		v1.EnvVar{
 			Name: "RSC_WSTUNNEL_SECRET",
 			ValueFrom: ptr.To(v1.EnvVarSource{

internal/remote/config/config.go

@@ -18,6 +18,8 @@ limitations under the License.
 package config
 
 import (
+	"fmt"
+
 	amaltheadevv1alpha1 "github.com/SwissDataScienceCenter/amalthea/api/v1alpha1"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -35,8 +37,11 @@ const (
 )
 
 const (
-	serverPortFlag = "server-port"
-	fakeStartFlag  = "fake-start"
+	serverPortFlag         = "server-port"
+	fakeStartFlag          = "fake-start"
+	sessionPortFlag        = "session-port"
+	sessionURLPathFlag     = "session-url-path"
+	readinessProbeTypeFlag = "readiness-probe-type"
 )
 
 type RemoteSessionControllerConfig struct {
@@ -53,6 +58,15 @@ type RemoteSessionControllerConfig struct {
 
 	// FakeStart if true, do not start the remote session and print debug information
 	FakeStart bool
+
+	// SessionPort is the port where the remote session is expected to be serving
+	SessionPort int32
+
+	// SessionURLPath is the URL path for the HTTP readiness probe
+	SessionURLPath string
+
+	// ReadinessProbeType is "none", "tcp", or "http"
+	ReadinessProbeType string
 }
 
 func SetFlags(cmd *cobra.Command) error {
@@ -72,6 +86,30 @@ func SetFlags(cmd *cobra.Command) error {
 		return err
 	}
 
+	cmd.Flags().Int32(sessionPortFlag, 0, "port the remote session is expected to be serving on")
+	if err := viper.BindPFlag(sessionPortFlag, cmd.Flags().Lookup(sessionPortFlag)); err != nil {
+		return err
+	}
+	if err := viper.BindEnv(sessionPortFlag, configUtils.AsEnvVarFlag(sessionPortFlag)); err != nil {
+		return err
+	}
+
+	cmd.Flags().String(sessionURLPathFlag, "/", "URL path for the HTTP readiness probe")
+	if err := viper.BindPFlag(sessionURLPathFlag, cmd.Flags().Lookup(sessionURLPathFlag)); err != nil {
+		return err
+	}
+	if err := viper.BindEnv(sessionURLPathFlag, configUtils.AsEnvVarFlag(sessionURLPathFlag)); err != nil {
+		return err
+	}
+
+	cmd.Flags().String(readinessProbeTypeFlag, "none", "readiness probe type: none, tcp, or http")
+	if err := viper.BindPFlag(readinessProbeTypeFlag, cmd.Flags().Lookup(readinessProbeTypeFlag)); err != nil {
+		return err
+	}
+	if err := viper.BindEnv(readinessProbeTypeFlag, configUtils.AsEnvVarFlag(readinessProbeTypeFlag)); err != nil {
+		return err
+	}
+
 	// Set up shared flags
 	if err := configUtils.SetFlags(cmd); err != nil {
 		return err
@@ -97,11 +135,20 @@ func GetConfig() (cfg RemoteSessionControllerConfig, err error) {
 
 	cfg.ServerPort = viper.GetInt32(serverPortFlag)
 	cfg.FakeStart = viper.GetBool(fakeStartFlag)
+	cfg.SessionPort = viper.GetInt32(sessionPortFlag)
+	cfg.SessionURLPath = viper.GetString(sessionURLPathFlag)
+	cfg.ReadinessProbeType = viper.GetString(readinessProbeTypeFlag)
 
 	return cfg, nil
 }
 
 func (cfg *RemoteSessionControllerConfig) Validate() error {
+	if cfg.ReadinessProbeType != string(amaltheadevv1alpha1.None) &&
+		cfg.ReadinessProbeType != string(amaltheadevv1alpha1.TCP) &&
+		cfg.ReadinessProbeType != string(amaltheadevv1alpha1.HTTP) {
+		return fmt.Errorf("invalid readiness probe type: %s", cfg.ReadinessProbeType)
+	}
+
 	// FireCREST has priority over Runai
 	cfg.RemoteKind = RemoteKindFirecrest
 	firecrestConfigErr := cfg.Firecrest.Validate()

internal/remote/config/config_test.go

@@ -19,6 +19,7 @@ package config
 import (
 	"testing"
 
+	amaltheadevv1alpha1 "github.com/SwissDataScienceCenter/amalthea/api/v1alpha1"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
@@ -86,3 +87,112 @@ func TestConfig(t *testing.T) {
 		})
 	}
 }
+
+func TestConfigNewFlags(t *testing.T) {
+	tests := []struct {
+		name            string
+		args            []string
+		env             map[string]string
+		wantProbeType   string
+		wantSessionPort int32
+		wantURLPath     string
+		wantErr         bool
+	}{
+		{
+			name: "default with empty readiness probe",
+			args: []string{
+				"--firecrest-api-url=https://firecrest.example.com",
+				"--firecrest-system-name=test-system",
+				"--auth-kind=client_credentials",
+				"--auth-token-uri=https://auth.example.com/token",
+				"--auth-firecrest-client-id=my-client",
+				"--auth-firecrest-client-secret=my-secret",
+			},
+			wantProbeType:   string(amaltheadevv1alpha1.None),
+			wantSessionPort: 0,
+			wantURLPath:     "/",
+		},
+		{
+			name: "custom flags via CLI",
+			args: []string{
+				"--readiness-probe-type=tcp",
+				"--session-port=8888",
+				"--session-url-path=/lab",
+				"--runai-base-url=https://runai.example.com",
+				"--runai-project=my-project",
+				"--auth-kind=client_credentials",
+				"--auth-token-uri=https://runai-auth.example.com/token",
+				"--auth-runai-client-id=runai-client",
+				"--auth-runai-client-secret=runai-secret",
+			},
+			wantProbeType:   string(amaltheadevv1alpha1.TCP),
+			wantSessionPort: 8888,
+			wantURLPath:     "/lab",
+		},
+		{
+			name: "custom flags via environment variables",
+			args: []string{
+				"--firecrest-api-url=https://firecrest.example.com",
+				"--firecrest-system-name=test-system",
+				"--auth-kind=client_credentials",
+				"--auth-token-uri=https://auth.example.com/token",
+				"--auth-firecrest-client-id=my-client",
+				"--auth-firecrest-client-secret=my-secret",
+			},
+			env: map[string]string{
+				"RSC_READINESS_PROBE_TYPE": "http",
+				"RSC_SESSION_PORT":         "9999",
+				"RSC_SESSION_URL_PATH":     "/jupyter",
+			},
+			wantProbeType:   string(amaltheadevv1alpha1.HTTP),
+			wantSessionPort: 9999,
+			wantURLPath:     "/jupyter",
+		},
+		{
+			name: "invalid readiness probe type errors",
+			args: []string{
+				"--readiness-probe-type=invalid",
+				"--firecrest-api-url=https://firecrest.example.com",
+				"--firecrest-system-name=test-system",
+				"--auth-kind=client_credentials",
+				"--auth-token-uri=https://auth.example.com/token",
+				"--auth-firecrest-client-id=my-client",
+				"--auth-firecrest-client-secret=my-secret",
+			},
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			viper.Reset()
+			for k, v := range tt.env {
+				t.Setenv(k, v)
+			}
+
+			cmd := &cobra.Command{
+				Use: "test",
+				Run: func(cmd *cobra.Command, args []string) {},
+			}
+			err := SetFlags(cmd)
+			require.NoError(t, err)
+
+			cmd.SetArgs(tt.args)
+			err = cmd.Execute()
+			require.NoError(t, err)
+
+			cfg, err := GetConfig()
+			require.NoError(t, err)
+
+			err = cfg.Validate()
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantProbeType, cfg.ReadinessProbeType)
+			assert.Equal(t, tt.wantSessionPort, cfg.SessionPort)
+			assert.Equal(t, tt.wantURLPath, cfg.SessionURLPath)
+		})
+	}
+}

internal/remote/server/server.go

@@ -20,11 +20,13 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"net"
 	"net/http"
 	"os"
 	"os/signal"
 	"time"
 
+	amaltheadevv1alpha1 "github.com/SwissDataScienceCenter/amalthea/api/v1alpha1"
 	"github.com/SwissDataScienceCenter/amalthea/internal/common"
 	"github.com/SwissDataScienceCenter/amalthea/internal/remote/config"
 	"github.com/SwissDataScienceCenter/amalthea/internal/remote/controller"
@@ -56,7 +58,7 @@ func Start() {
 		os.Exit(1)
 	}
 
-	server := newServer(controller)
+	server := newServer(controller, cfg)
 
 	address := fmt.Sprintf(":%d", cfg.ServerPort)
 
@@ -95,7 +97,7 @@ func Start() {
 var logLevel *slog.LevelVar = new(slog.LevelVar)
 var jsonLogger *slog.Logger = slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: logLevel}))
 
-func newServer(controller controller.RemoteSessionController) (server *echo.Echo) {
+func newServer(controller controller.RemoteSessionController, cfg config.RemoteSessionControllerConfig) (server *echo.Echo) {
 	e := echo.New()
 
 	e.HideBanner = true
@@ -114,7 +116,49 @@ func newServer(controller controller.RemoteSessionController) (server *echo.Echo
 
 	// Readiness endpoint
 	e.GET("/ready", func(c echo.Context) error {
-		return c.NoContent(http.StatusOK)
+		switch cfg.ReadinessProbeType {
+		case string(amaltheadevv1alpha1.TCP):
+			dialer := net.Dialer{}
+			dialCtx, dialCtxCancel := context.WithTimeout(c.Request().Context(), 5*time.Second)
+			defer dialCtxCancel()
+			conn, err := dialer.DialContext(dialCtx, "tcp", fmt.Sprintf("127.0.0.1:%d", cfg.SessionPort))
+
+			if err != nil {
+				return c.NoContent(http.StatusServiceUnavailable)
+			}
+			if err := conn.Close(); err != nil {
+				slog.Error("failed to close readiness probe connection", "error", err)
+			}
+			return c.NoContent(http.StatusOK)
+		case string(amaltheadevv1alpha1.HTTP):
+			client := &http.Client{
+				Timeout: 5 * time.Second,
+				CheckRedirect: func(req *http.Request, via []*http.Request) error {
+					return http.ErrUseLastResponse
+				},
+			}
+			url := fmt.Sprintf("http://127.0.0.1:%d%s", cfg.SessionPort, cfg.SessionURLPath)
+			req, err := http.NewRequestWithContext(c.Request().Context(), "GET", url, nil)
+			if err != nil {
+				return c.NoContent(http.StatusServiceUnavailable)
+			}
+			resp, err := client.Do(req)
+			if err != nil {
+				return c.NoContent(http.StatusServiceUnavailable)
+			}
+			defer func() {
+				if err := resp.Body.Close(); err != nil {
+					slog.Error("failed to close readiness probe response body", "error", err)
+				}
+			}()
+			if resp.StatusCode >= 200 && resp.StatusCode < 400 {
+				return c.NoContent(http.StatusOK)
+			}
+			return c.NoContent(http.StatusServiceUnavailable)
+		default:
+			// Case None or unset: preserve old behavior for backward compatibility
+			return c.NoContent(http.StatusOK)
+		}
 	})
 
 	// Status endpoint