Add server-lifecycle flags to ssh setup (--max-clients, --server-timeout)

.nextchanges/cli/ssh-setup-server-lifecycle-flags.md

@@ -0,0 +1 @@
+* Add `--max-clients` and `--server-timeout` flags to `databricks ssh setup`, and `--server-timeout` to `databricks ssh connect`. Both are fixed when the SSH tunnel server job is submitted, so `ssh setup` now serializes them into the generated `ProxyCommand` instead of falling back to the built-in defaults. ([#6547](https://github.com/databricks/cli/pull/6547))

acceptance/ssh/setup/output.txt

@@ -0,0 +1,18 @@
+
+=== ProxyCommand written by setup --max-clients=25 --server-timeout=48h
+ssh connect --proxy --cluster=[TEST_DEFAULT_CLUSTER_ID] --auto-start-cluster=true --shutdown-delay=10m0s --max-clients=25 --server-timeout=48h0m0s
+
+=== A shutdown delay beyond the default lifetime raises it, no --server-timeout needed
+ssh connect --proxy --cluster=[TEST_DEFAULT_CLUSTER_ID] --auto-start-cluster=true --shutdown-delay=48h0m0s --max-clients=10 --server-timeout=48h0m0s
+
+=== Rejects a server that would refuse every connection
+>>> [CLI] ssh setup --name=broken --cluster=[TEST_DEFAULT_CLUSTER_ID] --max-clients=0
+Error: --max-clients must be at least 1, got 0
+
+=== Rejects a shutdown delay the server can never reach
+>>> [CLI] ssh setup --name=broken --cluster=[TEST_DEFAULT_CLUSTER_ID] --shutdown-delay=48h --server-timeout=24h
+Error: --shutdown-delay (48h0m0s) cannot be longer than --server-timeout (24h0m0s)
+
+=== No host config is written for the rejected setups
+home/.databricks/ssh-tunnel-configs/long-delay
+home/.databricks/ssh-tunnel-configs/my-cluster

acceptance/ssh/setup/script

@@ -0,0 +1,27 @@
+sethome "./home"
+
+# Send the command's own output to a LOG file: it echoes the whole host config, whose IdentityFile
+# and ProxyCommand carry OS-specific absolute paths.
+$CLI ssh setup --name=my-cluster --cluster=$TEST_DEFAULT_CLUSTER_ID --max-clients=25 --server-timeout=48h &>LOG.setup
+
+# The ProxyCommand is the invocation that submits the SSH server job, so --max-clients and
+# --server-timeout have to appear there or `ssh my-cluster` silently starts a server with the
+# defaults.
+title "ProxyCommand written by setup --max-clients=25 --server-timeout=48h\n"
+sed -n 's/.*\(ssh connect --proxy.*\)/\1/p' "$HOME/.databricks/ssh-tunnel-configs/my-cluster"
+
+# A shutdown delay longer than the default 24h lifetime worked before --server-timeout existed,
+# and `ssh setup` persisted it into the ProxyCommand. OpenSSH runs that ProxyCommand verbatim, so
+# the delay has to keep raising the lifetime rather than becoming an error the user cannot fix.
+title "A shutdown delay beyond the default lifetime raises it, no --server-timeout needed\n"
+$CLI ssh setup --name=long-delay --cluster=$TEST_DEFAULT_CLUSTER_ID --shutdown-delay=48h &>LOG.long-delay
+sed -n 's/.*\(ssh connect --proxy.*\)/\1/p' "$HOME/.databricks/ssh-tunnel-configs/long-delay"
+
+title "Rejects a server that would refuse every connection"
+musterr trace $CLI ssh setup --name=broken --cluster=$TEST_DEFAULT_CLUSTER_ID --max-clients=0
+
+title "Rejects a shutdown delay the server can never reach"
+musterr trace $CLI ssh setup --name=broken --cluster=$TEST_DEFAULT_CLUSTER_ID --shutdown-delay=48h --server-timeout=24h
+
+title "No host config is written for the rejected setups\n"
+find.py 'ssh-tunnel-configs' --expect 2

acceptance/ssh/setup/test.toml

@@ -0,0 +1,8 @@
+# `ssh setup` only reads the cluster and writes local SSH config, so a real workspace adds nothing.
+Cloud = false
+
+Ignore = [
+  "home",
+]
+
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

experimental/ssh/cmd/connect.go

@@ -7,8 +7,25 @@ import (
 	"github.com/databricks/cli/experimental/ssh/internal/client"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 )
 
+// resolveServerTimeout returns the lifetime to submit the SSH tunnel server job with.
+//
+// Before --server-timeout existed the lifetime was max(24h, --shutdown-delay), so
+// `ssh setup --shutdown-delay=48h` produced a working 48h tunnel and persisted that delay into
+// the generated ProxyCommand. Keep honoring it whenever --server-timeout is not set: OpenSSH
+// runs a persisted ProxyCommand verbatim, so a user whose host config predates the flag cannot
+// add it, and rejecting the pair would break `ssh <name>` with an error naming a flag they have
+// no way to pass. An explicit --server-timeout always wins, and ClientOptions.Validate still
+// rejects a shutdown delay that outlives a lifetime the user asked for explicitly.
+func resolveServerTimeout(flags *pflag.FlagSet, serverTimeout, shutdownDelay time.Duration) time.Duration {
+	if flags.Changed("server-timeout") {
+		return serverTimeout
+	}
+	return max(serverTimeout, shutdownDelay)
+}
+
 func newConnectCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "connect",
@@ -32,6 +49,7 @@ Connect to a dedicated cluster:
 	var serverMetadata string
 	var shutdownDelay time.Duration
 	var maxClients int
+	var serverTimeout time.Duration
 	var handoverTimeout time.Duration
 	var releasesDir string
 	var autoStartCluster bool
@@ -46,6 +64,7 @@ Connect to a dedicated cluster:
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks dedicated cluster ID")
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server after the last client disconnects")
 	cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients")
+	cmd.Flags().DurationVar(&serverTimeout, "server-timeout", defaultServerTimeout, "Maximum lifetime of the SSH server; it is terminated after this duration even if clients are connected")
 	cmd.Flags().BoolVar(&autoStartCluster, "auto-start-cluster", true, "Automatically start the cluster if it is not running")
 
 	cmd.Flags().StringVar(&connectionName, "name", "", "Connection name to reuse across sessions (serverless only)")
@@ -121,7 +140,7 @@ Connect to a dedicated cluster:
 			HandoverTimeout:      handoverTimeout,
 			KeepaliveInterval:    defaultKeepaliveInterval,
 			ReleasesDir:          releasesDir,
-			ServerTimeout:        max(serverTimeout, shutdownDelay),
+			ServerTimeout:        resolveServerTimeout(cmd.Flags(), serverTimeout, shutdownDelay),
 			TaskStartupTimeout:   startupTimeout,
 			AutoStartCluster:     autoStartCluster,
 			ClientPublicKeyName:  clientPublicKeyName,

experimental/ssh/cmd/constants.go

@@ -13,8 +13,10 @@ const (
 	// 30 second SSH-level keepalive that was verified to prevent it.
 	defaultKeepaliveInterval  = 20 * time.Second
 	defaultEnvironmentVersion = 4
+	// Default cap on how long an SSH tunnel server is allowed to live. Fixed when the
+	// server job is submitted, so it is only settable by the invocation that starts it.
+	defaultServerTimeout = 24 * time.Hour
 
-	serverTimeout         = 24 * time.Hour
 	taskStartupTimeout    = 10 * time.Minute
 	gpuTaskStartupTimeout = 45 * time.Minute
 	serverPortRange       = 100

experimental/ssh/cmd/setup.go

@@ -24,6 +24,8 @@ For serverless connections, use ` + "`databricks ssh connect`" + ` (no setup ste
 	var clusterID string
 	var sshConfigPath string
 	var shutdownDelay time.Duration
+	var maxClients int
+	var serverTimeout time.Duration
 	var autoStartCluster bool
 	var autoApprove bool
 
@@ -33,6 +35,8 @@ For serverless connections, use ` + "`databricks ssh connect`" + ` (no setup ste
 	cmd.Flags().BoolVar(&autoStartCluster, "auto-start-cluster", true, "Automatically start the cluster when establishing the ssh connection")
 	cmd.Flags().StringVar(&sshConfigPath, "ssh-config", "", "Path to SSH config file (default ~/.ssh/config)")
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "SSH server will terminate after this delay if there are no active connections")
+	cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients")
+	cmd.Flags().DurationVar(&serverTimeout, "server-timeout", defaultServerTimeout, "Maximum lifetime of the SSH server; it is terminated after this duration even if clients are connected")
 	cmd.Flags().BoolVar(&autoApprove, "auto-approve", false, "Skip confirmation prompts, recreating existing SSH host configs without asking")
 
 	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
@@ -51,6 +55,8 @@ For serverless connections, use ` + "`databricks ssh connect`" + ` (no setup ste
 			AutoStartCluster: autoStartCluster,
 			SSHConfigPath:    sshConfigPath,
 			ShutdownDelay:    shutdownDelay,
+			MaxClients:       maxClients,
+			ServerTimeout:    resolveServerTimeout(cmd.Flags(), serverTimeout, shutdownDelay),
 			Profile:          wsClient.Config.Profile,
 			AutoApprove:      autoApprove,
 		}

experimental/ssh/cmd/setup_test.go

@@ -0,0 +1,183 @@
+package ssh
+
+import (
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/databricks/cli/experimental/ssh/internal/client"
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// The ProxyCommand written by `ssh setup` is what OpenSSH later executes, and it is the
+// invocation that submits the SSH server job. So every flag setup serializes has to exist on
+// `ssh connect` and survive the round trip - otherwise the value the user passed to setup is
+// silently replaced by connect's own default.
+func TestSetupProxyCommandRoundTripsThroughConnect(t *testing.T) {
+	opts := client.ClientOptions{
+		ClusterID:        "cluster-123",
+		AutoStartCluster: true,
+		ShutdownDelay:    15 * time.Minute,
+		MaxClients:       25,
+		ServerTimeout:    48 * time.Hour,
+	}
+	proxyCommand, err := opts.ToProxyCommand()
+	require.NoError(t, err)
+
+	_, args, found := strings.Cut(proxyCommand, " ssh connect ")
+	require.True(t, found, "proxy command %q must invoke 'ssh connect'", proxyCommand)
+
+	flags := newConnectCommand().Flags()
+	require.NoError(t, flags.Parse(strings.Fields(args)))
+
+	proxyMode, err := flags.GetBool("proxy")
+	require.NoError(t, err)
+	assert.True(t, proxyMode)
+
+	clusterID, err := flags.GetString("cluster")
+	require.NoError(t, err)
+	assert.Equal(t, opts.ClusterID, clusterID)
+
+	autoStart, err := flags.GetBool("auto-start-cluster")
+	require.NoError(t, err)
+	assert.Equal(t, opts.AutoStartCluster, autoStart)
+
+	shutdownDelay, err := flags.GetDuration("shutdown-delay")
+	require.NoError(t, err)
+	assert.Equal(t, opts.ShutdownDelay, shutdownDelay)
+
+	maxClients, err := flags.GetInt("max-clients")
+	require.NoError(t, err)
+	assert.Equal(t, opts.MaxClients, maxClients)
+
+	serverTimeout, err := flags.GetDuration("server-timeout")
+	require.NoError(t, err)
+	assert.Equal(t, opts.ServerTimeout, serverTimeout)
+}
+
+// Both commands start the same server, so a value the user does not override has to mean the
+// same thing whether the tunnel was configured with setup or started with connect.
+func TestSetupAndConnectShareServerLifecycleDefaults(t *testing.T) {
+	setupFlags := newSetupCommand().Flags()
+	connectFlags := newConnectCommand().Flags()
+
+	for _, name := range []string{"shutdown-delay", "max-clients", "server-timeout"} {
+		setupFlag := setupFlags.Lookup(name)
+		connectFlag := connectFlags.Lookup(name)
+		require.NotNil(t, setupFlag, "setup is missing --%s", name)
+		require.NotNil(t, connectFlag, "connect is missing --%s", name)
+		assert.Equal(t, connectFlag.DefValue, setupFlag.DefValue, "--%s default differs between setup and connect", name)
+	}
+}
+
+// A ProxyCommand written by `ssh setup --shutdown-delay=48h` before --server-timeout existed
+// carries no --server-timeout, and OpenSSH runs it verbatim - so the delay has to keep raising
+// the server's lifetime the way it did then, or `ssh <name>` breaks for a host config the user
+// cannot edit from the command line.
+func TestLegacyProxyCommandWithLongShutdownDelayStillWorks(t *testing.T) {
+	flags := newConnectCommand().Flags()
+	require.NoError(t, flags.Parse([]string{
+		"--proxy", "--cluster=abc-123", "--auto-start-cluster=true", "--shutdown-delay=48h0m0s",
+	}))
+
+	shutdownDelay, err := flags.GetDuration("shutdown-delay")
+	require.NoError(t, err)
+	serverTimeout, err := flags.GetDuration("server-timeout")
+	require.NoError(t, err)
+
+	opts := client.ClientOptions{
+		ProxyMode:     true,
+		ClusterID:     "abc-123",
+		MaxClients:    defaultMaxClients,
+		ShutdownDelay: shutdownDelay,
+		ServerTimeout: resolveServerTimeout(flags, serverTimeout, shutdownDelay),
+	}
+	require.NoError(t, opts.Validate())
+	assert.Equal(t, 48*time.Hour, opts.ServerTimeout, "the shutdown delay must raise the server's lifetime to cover it")
+}
+
+// The reason the max() above is gated: an explicitly requested lifetime must not be silently
+// widened by a longer shutdown delay. That pair is a real conflict, so it stays an error.
+func TestExplicitServerTimeoutShorterThanShutdownDelayIsRejected(t *testing.T) {
+	flags := newConnectCommand().Flags()
+	require.NoError(t, flags.Parse([]string{
+		"--proxy", "--cluster=abc-123", "--shutdown-delay=48h", "--server-timeout=24h",
+	}))
+
+	shutdownDelay, err := flags.GetDuration("shutdown-delay")
+	require.NoError(t, err)
+	serverTimeout, err := flags.GetDuration("server-timeout")
+	require.NoError(t, err)
+
+	resolved := resolveServerTimeout(flags, serverTimeout, shutdownDelay)
+	assert.Equal(t, 24*time.Hour, resolved, "an explicit --server-timeout wins over a longer --shutdown-delay")
+
+	opts := client.ClientOptions{
+		ProxyMode:     true,
+		ClusterID:     "abc-123",
+		MaxClients:    defaultMaxClients,
+		ShutdownDelay: shutdownDelay,
+		ServerTimeout: resolved,
+	}
+	assert.EqualError(t, opts.Validate(), "--shutdown-delay (48h0m0s) cannot be longer than --server-timeout (24h0m0s)")
+}
+
+func TestResolveServerTimeout(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+		want time.Duration
+	}{
+		{
+			name: "no flags: the default lifetime",
+			args: nil,
+			want: defaultServerTimeout,
+		},
+		{
+			name: "shutdown delay within the default lifetime leaves it alone",
+			args: []string{"--shutdown-delay=1h"},
+			want: defaultServerTimeout,
+		},
+		{
+			name: "shutdown delay beyond the default lifetime raises it",
+			args: []string{"--shutdown-delay=48h"},
+			want: 48 * time.Hour,
+		},
+		{
+			name: "an explicit lifetime is used as given",
+			args: []string{"--server-timeout=1h"},
+			want: time.Hour,
+		},
+		{
+			name: "an explicit lifetime is not widened by a longer shutdown delay",
+			args: []string{"--shutdown-delay=2h", "--server-timeout=1h"},
+			want: time.Hour,
+		},
+		{
+			// Same value as the default, but passed explicitly: the user asked for it, so a
+			// longer shutdown delay must not override it.
+			name: "an explicit lifetime equal to the default still wins",
+			args: []string{"--shutdown-delay=48h", "--server-timeout=24h"},
+			want: defaultServerTimeout,
+		},
+	}
+
+	// Both commands resolve the lifetime the same way, so check them together.
+	for _, newCmd := range []func() *cobra.Command{newSetupCommand, newConnectCommand} {
+		for _, tt := range tests {
+			t.Run(newCmd().Name()+"/"+tt.name, func(t *testing.T) {
+				flags := newCmd().Flags()
+				require.NoError(t, flags.Parse(tt.args))
+
+				shutdownDelay, err := flags.GetDuration("shutdown-delay")
+				require.NoError(t, err)
+				serverTimeout, err := flags.GetDuration("server-timeout")
+				require.NoError(t, err)
+
+				assert.Equal(t, tt.want, resolveServerTimeout(flags, serverTimeout, shutdownDelay))
+			})
+		}
+	}
+}

experimental/ssh/internal/client/client.go

@@ -155,6 +155,22 @@ func (o *ClientOptions) Validate() error {
 	if o.BaseEnvironment != "" && o.ClusterID != "" {
 		return errors.New("--base-environment can only be used with serverless compute")
 	}
+	// A server started with fewer than one client slot rejects every connection with an
+	// opaque websocket handshake failure, so catch it here instead.
+	if o.MaxClients < 1 {
+		return fmt.Errorf("--max-clients must be at least 1, got %d", o.MaxClients)
+	}
+	// The submitted job carries this as timeout_seconds, a whole number of seconds, and 0 means
+	// "no timeout" in the Jobs API. So every value below one second - not just zero - truncates
+	// to an unbounded run instead of the short-lived server that was asked for.
+	if o.ServerTimeout < time.Second {
+		return fmt.Errorf("--server-timeout must be at least 1s, got %s", o.ServerTimeout)
+	}
+	// The server only starts counting down the shutdown delay once the last client leaves, so a
+	// delay longer than the server's lifetime can never elapse.
+	if o.ShutdownDelay > o.ServerTimeout {
+		return fmt.Errorf("--shutdown-delay (%s) cannot be longer than --server-timeout (%s)", o.ShutdownDelay, o.ServerTimeout)
+	}
 	return nil
 }
 
@@ -229,6 +245,18 @@ func (o *ClientOptions) ToProxyCommand() (string, error) {
 			executablePath, o.ClusterID, o.AutoStartCluster, o.ShutdownDelay.String())
 	}
 
+	// Both of these are fixed when the server job is submitted, and for a host configured by
+	// `ssh setup` the submitting invocation is always the ProxyCommand, so they have to be
+	// carried here or the user's choice is lost. Zero means "not set": the receiving command
+	// then applies its own flag default.
+	if o.MaxClients > 0 {
+		proxyCommand += " --max-clients=" + strconv.Itoa(o.MaxClients)
+	}
+
+	if o.ServerTimeout > 0 {
+		proxyCommand += " --server-timeout=" + o.ServerTimeout.String()
+	}
+
 	if o.ServerMetadata != "" {
 		proxyCommand += " --metadata=" + o.ServerMetadata
 	}