Keep idle SSH tunnel sessions alive on dedicated clusters

.nextchanges/cli/ssh-tunnel-keepalive.md

@@ -1 +1 @@
-Fixed idle `databricks ssh connect` sessions disconnecting after a few minutes. The tunnel now sends a websocket keepalive every 20 seconds, so a session nobody is typing into stays connected without setting `ServerAliveInterval` in the SSH client config.
+Fixed idle `databricks ssh connect` sessions disconnecting after a few minutes, on dedicated clusters and on serverless. The tunnel now keeps itself warm: the SSH client and the SSH server on the compute exchange keepalives every 30 seconds, and the CLI's proxy pings the tunnel's websocket every 20 seconds. A session nobody is typing into stays connected, with no need to set `ServerAliveInterval` by hand.

experimental/ssh/internal/client/client.go

@@ -820,6 +820,7 @@ func buildSSHArgs(userName, privateKeyPath, proxyCommand, hostName, wsHome strin
 		"-o", "IdentitiesOnly=yes",
 		"-o", "StrictHostKeyChecking=accept-new",
 		"-o", "ConnectTimeout=360",
+		"-o", "ServerAliveInterval=" + strconv.Itoa(sshconfig.ServerAliveIntervalSeconds),
 		"-o", "ProxyCommand=" + proxyCommand,
 	}
 	if opts.UserKnownHostsFile != "" {

experimental/ssh/internal/client/client_internal_test.go

@@ -5,10 +5,13 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"slices"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/databricks/cli/experimental/ssh/internal/sshconfig"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/telemetry/protos"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
@@ -373,6 +376,17 @@ func TestBuildRemoteShellArgs(t *testing.T) {
 	})
 }
 
+func TestBuildSSHArgsSetsServerAliveInterval(t *testing.T) {
+	args := buildSSHArgs("user", "/key", "proxy command", "myhost", "", ClientOptions{})
+
+	// ssh stops parsing options at the destination, so an option placed after the host would be
+	// treated as part of the remote command rather than as an ssh option.
+	optIdx := slices.Index(args, "ServerAliveInterval="+strconv.Itoa(sshconfig.ServerAliveIntervalSeconds))
+	require.NotEqual(t, -1, optIdx, "ssh must be asked to send keepalives")
+	require.Equal(t, "-o", args[optIdx-1])
+	assert.Less(t, optIdx, slices.Index(args, "myhost"), "the option must precede the destination host")
+}
+
 func TestBuildSSHArgsPTYPlacement(t *testing.T) {
 	indexOf := func(args []string, want string) int {
 		for i, a := range args {

experimental/ssh/internal/server/sshd.go

@@ -9,6 +9,7 @@ import (
 	"os/exec"
 	"path"
 	"path/filepath"
+	"strconv"
 	"strings"
 
 	"github.com/databricks/cli/experimental/ssh/internal/keys"
@@ -17,6 +18,18 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 )
 
+// clientAliveIntervalSeconds is how often sshd asks the client to confirm it is still there. It
+// drives the keepalive from the server end of the tunnel; sshconfig.ServerAliveIntervalSeconds
+// documents why an SSH keepalive is what keeps the leg past the driver proxy from being reaped.
+// Configuring it here covers clients the CLI does not configure — a hand-written ProxyCommand host
+// block, or an IDE that supplies its own ssh options — where nothing sets ServerAliveInterval.
+// The two intervals are deliberately independent: neither package imports the other, and each end
+// keeps its own leg warm, so they need not track a single shared value.
+//
+// It also brings in ClientAliveCountMax (OpenSSH default 3), so sshd reclaims a session whose
+// client has gone away after ~90s instead of holding it open until the server's shutdown delay.
+const clientAliveIntervalSeconds = 30
+
 func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOptions) (string, error) {
 	clientPublicKey, err := keys.GetSecret(ctx, client, opts.SecretScopeName, opts.AuthorizedKeySecretName)
 	if err != nil {
@@ -76,15 +89,7 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient,
 	}
 	setEnv := setEnvBuf.String()
 
-	sshdConfigContent := "PubkeyAuthentication yes\n" +
-		"PasswordAuthentication no\n" +
-		"ChallengeResponseAuthentication no\n" +
-		"Subsystem sftp internal-sftp\n" +
-		"HostKey " + keyPath + "\n" +
-		"AuthorizedKeysFile " + authKeysPath + "\n" +
-		setEnv + "\n"
-
-	if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent), 0o600); err != nil {
+	if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent(keyPath, authKeysPath, setEnv)), 0o600); err != nil {
 		return "", err
 	}
 
@@ -97,6 +102,18 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient,
 	return sshdConfig, nil
 }
 
+// sshdConfigContent assembles the configuration the tunnel's sshd runs with.
+func sshdConfigContent(hostKeyPath, authorizedKeysPath, setEnv string) string {
+	return "PubkeyAuthentication yes\n" +
+		"PasswordAuthentication no\n" +
+		"ChallengeResponseAuthentication no\n" +
+		"ClientAliveInterval " + strconv.Itoa(clientAliveIntervalSeconds) + "\n" +
+		"Subsystem sftp internal-sftp\n" +
+		"HostKey " + hostKeyPath + "\n" +
+		"AuthorizedKeysFile " + authorizedKeysPath + "\n" +
+		setEnv + "\n"
+}
+
 func createSSHDProcess(ctx context.Context, configPath string) *exec.Cmd {
 	return exec.CommandContext(ctx, "/usr/sbin/sshd", "-f", configPath, "-i")
 }

experimental/ssh/internal/server/sshd_test.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -71,3 +72,18 @@ func TestEscapeEnvValue(t *testing.T) {
 		})
 	}
 }
+
+func TestSSHDConfigSetsClientAliveInterval(t *testing.T) {
+	config := sshdConfigContent("/keys/server-private-key", "/keys/authorized_keys", `SetEnv FOO="bar"`)
+
+	// This is the half of the keepalive that reaches clients the CLI never configures, so sshd has
+	// to drive it: without ClientAliveInterval sshd sends nothing on an idle session.
+	assert.Contains(t, config, "\nClientAliveInterval "+strconv.Itoa(clientAliveIntervalSeconds)+"\n")
+
+	// The interval has to fire well inside the ~8 minute reap window, and ClientAliveCountMax
+	// (OpenSSH default 3) intervals have to outlast the longest legitimate pause on a healthy
+	// tunnel — the up to 30s a handover can hold the sending loop. The 30 below is
+	// proxy.proxyHandoverInitTimeout's current value; it is unexported, so it can't be referenced.
+	assert.Less(t, clientAliveIntervalSeconds, 8*60)
+	assert.Greater(t, 3*clientAliveIntervalSeconds, 30)
+}

experimental/ssh/internal/sshconfig/sshconfig.go

@@ -19,6 +19,18 @@ const (
 	configDirName = ".databricks/ssh-tunnel-configs"
 )
 
+// ServerAliveIntervalSeconds is how often the ssh client asks the SSH server to confirm it is
+// still there. The reply is a real SSH packet, so the keepalive puts payload bytes on every hop
+// of the tunnel — and payload is what an idle session needs. The driver proxy terminates
+// websocket control frames itself, so the proxy's own websocket ping never becomes payload on
+// the leg past it, and that leg is reaped after ~8 minutes without any.
+//
+// It also brings in ServerAliveCountMax (OpenSSH default 3), so a tunnel that stops responding
+// is torn down after ~90s with ssh's own "server not responding" message instead of hanging.
+// That is well clear of the up to 30s a handover can hold the sending loop
+// (proxyHandoverInitTimeout), the longest legitimate pause on a healthy tunnel.
+const ServerAliveIntervalSeconds = 30
+
 func GetConfigDir(ctx context.Context) (string, error) {
 	homeDir, err := env.UserHomeDir(ctx)
 	if err != nil {
@@ -206,9 +218,10 @@ func GenerateHostConfig(hostName, userName, identityFile, proxyCommand string) s
 Host %s
     User %s
     ConnectTimeout 360
+    ServerAliveInterval %d
     StrictHostKeyChecking accept-new
     IdentitiesOnly yes
     IdentityFile %q
     ProxyCommand %s
-`, hostName, userName, identityFile, proxyCommand)
+`, hostName, userName, ServerAliveIntervalSeconds, identityFile, proxyCommand)
 }

experimental/ssh/internal/sshconfig/sshconfig_test.go

@@ -1,6 +1,7 @@
 package sshconfig
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
@@ -10,6 +11,21 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestGenerateHostConfigSetsServerAliveInterval(t *testing.T) {
+	config := GenerateHostConfig("myhost", "root", "/keys/myhost", "databricks ssh connect --proxy")
+
+	// `ssh setup` and `--ide` reach ssh through this block and nothing else, so the option has to
+	// be in it.
+	assert.Contains(t, config, fmt.Sprintf("\n    ServerAliveInterval %d\n", ServerAliveIntervalSeconds))
+
+	// The interval has to fire well inside the ~8 minute reap window, and ServerAliveCountMax
+	// (OpenSSH default 3) intervals have to outlast the longest legitimate pause on a healthy
+	// tunnel — the up to 30s a handover can hold the sending loop. The 30 below is
+	// proxy.proxyHandoverInitTimeout's current value; it is unexported, so it can't be referenced.
+	assert.Less(t, ServerAliveIntervalSeconds, 8*60)
+	assert.Greater(t, 3*ServerAliveIntervalSeconds, 30)
+}
+
 func TestGetConfigDir(t *testing.T) {
 	dir, err := GetConfigDir(t.Context())
 	assert.NoError(t, err)