From 1ec34c9a795b25a50a64f0d973e95437626d1d2a Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 24 Aug 2026 08:24:09 +0000
Subject: [PATCH 1/3] Keep idle SSH tunnel sessions alive with a websocket ping
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

An idle `databricks ssh connect` session dies after roughly nine minutes.
Both proxy loops are purely data-driven, so a session nobody is typing into
puts no frames on the websocket at all, and the server side reaps the stream
it then considers dead (websocket close 4000, Armeria ClosedStreamException).
Setting `ServerAliveInterval` in the SSH client config works around it
entirely, because SSH-level keepalives are real payload bytes that the sending
loop forwards — which is what pinned the diagnosis on the transport.

The client proxy now pings the websocket every 20 seconds for the life of the
connection, as an additional goroutine in the errgroup that already drives the
periodic handover tick. That placement makes the keepalive client-only by
construction: the server never calls RunClientProxy, so no flag can enable
server-side pinging, and the cluster-side binary needs no redeployment.

Pings take the proxy's existing serialised write path (sendMessage), which
already holds the handover mutex, so the serialisation gorilla/websocket
requires is inherited rather than newly built. A ping that ticks during a
handover blocks and goes out late; a handover establishes a fresh connection,
so the peer's idle clock resets anyway. A failed ping is logged at debug level
and never returned: the receiving loop stays the sole authority on whether the
connection is dead, and an error here would cancel the session the keepalive
exists to preserve. Liveness posture is keep-warm only — no read or write
deadlines, and the pong handler only logs.

The tunnel already rotates its websocket on a schedule via the periodic
handover, so hostility to long-lived streams on this transport was already
known and designed around here; the keepalive is the missing half of that
story.

Also corrects the `ssh server --shutdown-delay` help text, which claimed the
server shuts down "after no pings from clients" when no pings existed anywhere
in the tunnel — inaccurate today, and actively misleading once real pings exist.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 .nextchanges/cli/ssh-tunnel-keepalive.md      |   1 +
 experimental/ssh/cmd/connect.go               |   1 +
 experimental/ssh/cmd/constants.go             |  13 +-
 experimental/ssh/cmd/server.go                |   2 +-
 experimental/ssh/internal/client/client.go    |   4 +-
 experimental/ssh/internal/proxy/client.go     |  48 ++++-
 .../ssh/internal/proxy/client_server_test.go  |  39 ++--
 .../ssh/internal/proxy/keepalive_test.go      | 173 ++++++++++++++++++
 experimental/ssh/internal/proxy/proxy_test.go |  14 ++
 9 files changed, 271 insertions(+), 24 deletions(-)
 create mode 100644 .nextchanges/cli/ssh-tunnel-keepalive.md
 create mode 100644 experimental/ssh/internal/proxy/keepalive_test.go

diff --git a/.nextchanges/cli/ssh-tunnel-keepalive.md b/.nextchanges/cli/ssh-tunnel-keepalive.md
new file mode 100644
index 00000000000..71802b97adb
--- /dev/null
+++ b/.nextchanges/cli/ssh-tunnel-keepalive.md
@@ -0,0 +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.
diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go
index 2c50b871902..7502d70f560 100644
--- a/experimental/ssh/cmd/connect.go
+++ b/experimental/ssh/cmd/connect.go
@@ -119,6 +119,7 @@ Connect to a dedicated cluster:
 			ShutdownDelay:        shutdownDelay,
 			MaxClients:           maxClients,
 			HandoverTimeout:      handoverTimeout,
+			KeepaliveInterval:    defaultKeepaliveInterval,
 			ReleasesDir:          releasesDir,
 			ServerTimeout:        max(serverTimeout, shutdownDelay),
 			TaskStartupTimeout:   startupTimeout,
diff --git a/experimental/ssh/cmd/constants.go b/experimental/ssh/cmd/constants.go
index 64c99b5bd48..dd5d3b2fdf2 100644
--- a/experimental/ssh/cmd/constants.go
+++ b/experimental/ssh/cmd/constants.go
@@ -3,10 +3,15 @@ package ssh
 import "time"
 
 const (
-	defaultServerPort         = 7772
-	defaultMaxClients         = 10
-	defaultShutdownDelay      = 10 * time.Minute
-	defaultHandoverTimeout    = 30 * time.Minute
+	defaultServerPort      = 7772
+	defaultMaxClients      = 10
+	defaultShutdownDelay   = 10 * time.Minute
+	defaultHandoverTimeout = 30 * time.Minute
+	// How often the client pings the tunnel websocket so an idle SSH session keeps the transport
+	// alive. Matches the keepalive interval of the vite bridge (libs/apps/vite/bridge.go), and sits
+	// well under both the ~9 minutes after which idle sessions were observed to drop and the
+	// 30 second SSH-level keepalive that was verified to prevent it.
+	defaultKeepaliveInterval  = 20 * time.Second
 	defaultEnvironmentVersion = 4
 
 	serverTimeout         = 24 * time.Hour
diff --git a/experimental/ssh/cmd/server.go b/experimental/ssh/cmd/server.go
index 21c651b2365..3675a4a7fe8 100644
--- a/experimental/ssh/cmd/server.go
+++ b/experimental/ssh/cmd/server.go
@@ -41,7 +41,7 @@ and proxies them to local SSH daemon processes.`,
 	cmd.MarkFlagRequired("authorized-key-secret-name")
 
 	cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients")
-	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down after no pings from clients")
+	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server when there are no active connections")
 	cmd.Flags().StringVar(&version, "version", "", "Client version of the Databricks CLI")
 	cmd.Flags().BoolVar(&serverless, "serverless", false, "Enable serverless mode for Jupyter initialization")
 	cmd.Flags().StringVar(&usagePolicyID, "usage-policy-id", "", "Usage policy ID the job was submitted with")
diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go
index 90e99f6b590..b2a8c9d0b7f 100644
--- a/experimental/ssh/internal/client/client.go
+++ b/experimental/ssh/internal/client/client.go
@@ -86,6 +86,8 @@ type ClientOptions struct {
 	ServerMetadata string
 	// How often the CLI should reconnect to the server with new auth.
 	HandoverTimeout time.Duration
+	// How often the CLI pings the tunnel websocket to keep an idle session alive.
+	KeepaliveInterval time.Duration
 	// Max amount of time the server process is allowed to live
 	ServerTimeout time.Duration
 	// Max amount of time to wait for the SSH server task to reach RUNNING state
@@ -895,7 +897,7 @@ func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, server
 	requestHandoverTick := func() <-chan time.Time {
 		return time.After(opts.HandoverTimeout)
 	}
-	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, createConn)
+	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, createConn)
 }
 
 // accessModeUILabel maps a cluster's access mode to the name shown in the Databricks UI.
diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go
index a1e8389e7ff..47c4cb09b6d 100644
--- a/experimental/ssh/internal/proxy/client.go
+++ b/experimental/ssh/internal/proxy/client.go
@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/libs/log"
+	"github.com/gorilla/websocket"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -37,8 +38,25 @@ func (f *firstByteWriter) Write(p []byte) (int, error) {
 	return f.w.Write(p)
 }
 
-func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, createConn createWebsocketConnectionFunc) error {
-	proxy := newProxyConnection(createConn)
+// logPongs wraps a connection factory so every connection it creates — the initial one and each
+// one a handover creates — logs the pongs coming back for our keepalive pings. Debug visibility only:
+// the receiving loop stays the only judge of whether a connection is alive.
+func logPongs(ctx context.Context, createConn createWebsocketConnectionFunc) createWebsocketConnectionFunc {
+	return func(connCtx context.Context, connID string) (*websocket.Conn, error) {
+		conn, err := createConn(connCtx, connID)
+		if err != nil {
+			return nil, err
+		}
+		conn.SetPongHandler(func(string) error {
+			log.Debugf(ctx, "Received websocket keepalive pong")
+			return nil
+		})
+		return conn, nil
+	}
+}
+
+func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, createConn createWebsocketConnectionFunc) error {
+	proxy := newProxyConnection(logPongs(ctx, createConn))
 	log.Infof(ctx, "Establishing SSH proxy connection...")
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
@@ -68,9 +86,33 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque
 				}
 			}
 		})
+		g.Go(func() error {
+			// Keep the websocket carrying traffic while the SSH session is idle. Both proxy loops
+			// are data-driven, so an idle session puts no frames on the connection at all and the
+			// server side reaps the stream it then considers dead (websocket close 4000).
+			ticker := time.NewTicker(keepaliveInterval)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-gCtx.Done():
+					return gCtx.Err()
+				case <-ticker.C:
+					// Pings take the same serialised write path as data — gorilla forbids
+					// concurrent writers — so a ping that ticks during a handover blocks until it
+					// finishes and then goes out late. Harmless: a handover establishes a fresh
+					// connection, which resets the peer's idle clock anyway.
+					if err := proxy.sendMessage(websocket.PingMessage, nil); err != nil {
+						// Never fatal. A failed ping knows nothing the data loops don't, and an
+						// error returned here would cancel the session it exists to preserve.
+						// The receiving loop notices a genuinely dead connection within one read.
+						log.Debugf(gCtx, "Failed to send websocket keepalive ping: %v", err)
+					}
+				}
+			}
+		})
 		g.Go(func() error {
 			// When proxy.start returns (EOF from ssh, or the server closing the connection),
-			// cancel so the handover goroutine stops too and g.Wait can return.
+			// cancel so the handover and keepalive goroutines stop too and g.Wait can return.
 			defer cancel()
 			return proxy.start(gCtx, src, wrappedDst)
 		})
diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go
index 1915cc07c88..a52bc1a919d 100644
--- a/experimental/ssh/internal/proxy/client_server_test.go
+++ b/experimental/ssh/internal/proxy/client_server_test.go
@@ -38,7 +38,7 @@ type testClient struct {
 	Cleanup     func()
 }
 
-func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, errChan chan error) *testClient {
+func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, errChan chan error) *testClient {
 	ctx := cmdio.MockDiscard(t.Context())
 	clientInput, clientInputWriter := io.Pipe()
 	clientOutput := newTestBuffer(t)
@@ -49,13 +49,11 @@ func createTestClient(t *testing.T, serverURL string, requestHandoverTick func()
 		return conn, err
 	}
 	if requestHandoverTick == nil {
-		requestHandoverTick = func() <-chan time.Time {
-			return time.After(time.Hour)
-		}
+		requestHandoverTick = neverTick
 	}
 	wg := sync.WaitGroup{}
 	wg.Go(func() {
-		err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, createConn)
+		err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, keepaliveInterval, createConn)
 		if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, io.ErrClosedPipe) {
 			if errChan != nil {
 				errChan <- err
@@ -78,7 +76,7 @@ func createTestClient(t *testing.T, serverURL string, requestHandoverTick func()
 func TestClientServerEcho(t *testing.T) {
 	server := createTestServer(t, 2, time.Hour)
 	defer server.Close()
-	client := createTestClient(t, server.URL, nil, nil)
+	client := createTestClient(t, server.URL, nil, time.Hour, nil)
 	defer client.Cleanup()
 
 	testMsg1 := []byte("test message 1\n")
@@ -100,9 +98,9 @@ func TestClientServerEcho(t *testing.T) {
 func TestMultipleClients(t *testing.T) {
 	server := createTestServer(t, 2, time.Hour)
 	defer server.Close()
-	client1 := createTestClient(t, server.URL, nil, nil)
+	client1 := createTestClient(t, server.URL, nil, time.Hour, nil)
 	defer client1.Cleanup()
-	client2 := createTestClient(t, server.URL, nil, nil)
+	client2 := createTestClient(t, server.URL, nil, time.Hour, nil)
 	defer client2.Cleanup()
 
 	messageCount := 10
@@ -131,9 +129,9 @@ func TestMaxClients(t *testing.T) {
 	maxClients := 2
 	server := createTestServer(t, maxClients, time.Hour)
 	defer server.Close()
-	client1 := createTestClient(t, server.URL, nil, nil)
+	client1 := createTestClient(t, server.URL, nil, time.Hour, nil)
 	defer client1.Cleanup()
-	client2 := createTestClient(t, server.URL, nil, nil)
+	client2 := createTestClient(t, server.URL, nil, time.Hour, nil)
 	defer client2.Cleanup()
 
 	testMsg1 := []byte("test message 1\n")
@@ -147,7 +145,7 @@ func TestMaxClients(t *testing.T) {
 	require.NoError(t, err)
 
 	errChan := make(chan error, 1)
-	client3 := createTestClient(t, server.URL, nil, errChan)
+	client3 := createTestClient(t, server.URL, nil, time.Hour, errChan)
 	defer client3.Cleanup()
 	select {
 	case err = <-errChan:
@@ -158,6 +156,17 @@ func TestMaxClients(t *testing.T) {
 }
 
 func TestHandover(t *testing.T) {
+	t.Run("without keepalive", func(t *testing.T) {
+		runHandoverExchange(t, time.Hour)
+	})
+	// Pings share the proxy's serialised write path with the data stream: they must not corrupt or
+	// reorder it, nor trip gorilla's concurrent-write panic.
+	t.Run("with keepalive", func(t *testing.T) {
+		runHandoverExchange(t, time.Millisecond)
+	})
+}
+
+func runHandoverExchange(t *testing.T, keepaliveInterval time.Duration) {
 	server := createTestServer(t, 2, time.Hour)
 	defer server.Close()
 
@@ -165,7 +174,7 @@ func TestHandover(t *testing.T) {
 	requestHandoverTick := func() <-chan time.Time {
 		return handoverChan
 	}
-	client := createTestClient(t, server.URL, requestHandoverTick, nil)
+	client := createTestClient(t, server.URL, requestHandoverTick, keepaliveInterval, nil)
 	defer client.Cleanup()
 
 	var expectedOutput []byte
@@ -204,7 +213,7 @@ func TestQuickHandover(t *testing.T) {
 	requestHandoverTick := func() <-chan time.Time {
 		return handoverChan
 	}
-	client := createTestClient(t, server.URL, requestHandoverTick, nil)
+	client := createTestClient(t, server.URL, requestHandoverTick, time.Hour, nil)
 	defer client.Cleanup()
 
 	var expectedOutput []byte
@@ -256,7 +265,7 @@ func TestClientExitsWhenServerCommandFails(t *testing.T) {
 
 	done := make(chan error, 1)
 	go func() {
-		done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, createConn)
+		done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn)
 	}()
 
 	select {
@@ -303,7 +312,7 @@ func TestClientTimesOutWhenServerSendsNothing(t *testing.T) {
 
 	done := make(chan error, 1)
 	go func() {
-		done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, createConn)
+		done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn)
 	}()
 
 	select {
diff --git a/experimental/ssh/internal/proxy/keepalive_test.go b/experimental/ssh/internal/proxy/keepalive_test.go
new file mode 100644
index 00000000000..d1d779d4bff
--- /dev/null
+++ b/experimental/ssh/internal/proxy/keepalive_test.go
@@ -0,0 +1,173 @@
+package proxy
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+// startKeepaliveTestServer stands up a websocket peer that behaves like the SSH server side of the
+// tunnel for an idle session: it sends the first bytes (which the client waits for before it
+// considers the session established), then only reads. Pings it receives are reported on the
+// returned channel and answered with a pong, the same way gorilla's default ping handler does.
+func startKeepaliveTestServer(t *testing.T) (*httptest.Server, <-chan struct{}) {
+	pings := make(chan struct{}, 1)
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+		conn.SetPingHandler(func(appData string) error {
+			select {
+			case pings <- struct{}{}:
+			default:
+			}
+			return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second))
+		})
+		if err := conn.WriteMessage(websocket.BinaryMessage, []byte("SSH-2.0-test\r\n")); err != nil {
+			return
+		}
+		// Ping handlers only run while a read is in progress, so keep reading until the client goes away.
+		for {
+			if _, _, err := conn.ReadMessage(); err != nil {
+				return
+			}
+		}
+	}))
+	return server, pings
+}
+
+func keepaliveTestDialer(serverURL string, onConn func(*websocket.Conn)) createWebsocketConnectionFunc {
+	wsURL := "ws" + serverURL[4:]
+	return func(ctx context.Context, connID string) (*websocket.Conn, error) {
+		conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose
+		if err != nil {
+			return nil, err
+		}
+		if onConn != nil {
+			onConn(conn)
+		}
+		return conn, nil
+	}
+}
+
+// TestKeepalivePingReachesServer covers the fix itself: an idle session sends no data, so without
+// the keepalive nothing at all crosses the websocket and the server side eventually reaps the
+// stream it considers dead.
+func TestKeepalivePingReachesServer(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
+	server, pings := startKeepaliveTestServer(t)
+	defer server.Close()
+
+	// Never written to: the session stays idle for the whole test.
+	src, _ := io.Pipe()
+	done := make(chan error, 1)
+	go func() {
+		done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, nil))
+	}()
+
+	select {
+	case <-pings:
+	case err := <-done:
+		t.Fatalf("session ended before a keepalive ping arrived: %v", err)
+	case <-time.After(10 * time.Second):
+		t.Fatal("no keepalive ping arrived at the server")
+	}
+}
+
+// TestKeepalivePingFailureDoesNotEndSession asserts a keepalive can never be the thing that ends a
+// session: it has no information the data loops lack, and it runs in the errgroup that drives them,
+// so a returned error would tear the session down.
+func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
+	server, _ := startKeepaliveTestServer(t)
+	defer server.Close()
+
+	failWrites := func(conn *websocket.Conn) {
+		// A deadline in the past fails every write on this connection, so every ping fails.
+		// Set before the connection is handed to the proxy, so no writer can be in flight.
+		// Reads are unaffected: the connection is otherwise healthy and the session must survive.
+		conn.SetWriteDeadline(time.Now().Add(-time.Second)) // nolint:errcheck
+	}
+
+	src, _ := io.Pipe()
+	done := make(chan error, 1)
+	go func() {
+		done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, failWrites))
+	}()
+
+	// Long enough for many pings to be attempted and fail.
+	select {
+	case err := <-done:
+		t.Fatalf("session ended after a failed keepalive ping: %v", err)
+	case <-time.After(2 * time.Second):
+	}
+}
+
+// TestKeepalivePingBlockedByHandoverDoesNotDeadlock asserts the invariant the keepalive design
+// rests on: a ping sent through the proxy's serialised write path blocks for the duration of a
+// handover, and a handover waits on the receiving loop rather than that write path, so the two
+// cannot deadlock each other.
+func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) {
+	ctx := t.Context()
+
+	server := setupTestServer(ctx, t)
+	defer server.Cleanup()
+
+	// Holds the handover open at the point where it has taken the write path but has not yet
+	// completed, which is when a ping must block rather than break the handover.
+	handoverDialing := make(chan struct{})
+	releaseHandover := make(chan struct{})
+	var dials atomic.Int32
+
+	client := setupTestClientWithDialHook(ctx, t, server.URL, func() {
+		if dials.Add(1) > 1 {
+			close(handoverDialing)
+			<-releaseHandover
+		}
+	})
+	defer client.Cleanup()
+
+	handoverDone := make(chan error, 1)
+	go func() {
+		handoverDone <- client.Proxy.initiateHandover(ctx)
+	}()
+	<-handoverDialing
+
+	pingDone := make(chan error, 1)
+	go func() {
+		pingDone <- client.Proxy.sendMessage(websocket.PingMessage, nil)
+	}()
+
+	select {
+	case err := <-pingDone:
+		t.Fatalf("ping was written while a handover held the write path: %v", err)
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	close(releaseHandover)
+
+	select {
+	case err := <-handoverDone:
+		require.NoError(t, err)
+	case <-time.After(10 * time.Second):
+		t.Fatal("handover deadlocked while a keepalive ping waited on the write path")
+	}
+	select {
+	case err := <-pingDone:
+		require.NoError(t, err)
+	case <-time.After(10 * time.Second):
+		t.Fatal("keepalive ping never completed after the handover finished")
+	}
+}
diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go
index 0e1db9021e2..1e3af34f40d 100644
--- a/experimental/ssh/internal/proxy/proxy_test.go
+++ b/experimental/ssh/internal/proxy/proxy_test.go
@@ -139,12 +139,26 @@ func createTestWebsocketConnection(url string) (*websocket.Conn, error) {
 	return conn, err
 }
 
+// neverTick is a tick channel that never fires, for tests that don't exercise a periodic behaviour.
+func neverTick() <-chan time.Time {
+	return time.After(time.Hour)
+}
+
 func setupTestClient(ctx context.Context, t *testing.T, serverURL string) *TestProxy {
+	return setupTestClientWithDialHook(ctx, t, serverURL, nil)
+}
+
+// setupTestClientWithDialHook is setupTestClient with a hook called on every websocket dial: the
+// initial connection and each one a handover creates.
+func setupTestClientWithDialHook(ctx context.Context, t *testing.T, serverURL string, onDial func()) *TestProxy {
 	ctx = log.NewContext(ctx, log.GetLogger(ctx).With("Client", true))
 	clientInput, clientInputWriter := io.Pipe()
 	clientOutput := newTestBuffer(t)
 	wsURL := "ws" + serverURL[4:]
 	clientProxy := newProxyConnection(func(ctx context.Context, connID string) (*websocket.Conn, error) {
+		if onDial != nil {
+			onDial()
+		}
 		return createTestWebsocketConnection(wsURL)
 	})
 	err := clientProxy.connect(ctx)

From 81d83b890e32b79f58b6a14cefc37793324d69ee Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 24 Aug 2026 09:42:00 +0000
Subject: [PATCH 2/3] Log each keepalive ping, not just failures
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

End-to-end verification against dogfood showed the far end never returns a
pong: across 21-minute, 45-minute and forced-handover runs the pong handler
logged nothing, while ping writes never failed. Control frames do not make the
round trip on this transport, and the outbound ping alone is what keeps the
stream from being reaped.

That leaves a support engineer reading a customer's debug log with no positive
evidence that keepalives were flowing — only the absence of failures, which is
indistinguishable from a build without the keepalive. Log each successful ping
instead, at debug level: three lines a minute on a transport whose debug log
already carries full HTTP bodies.

Verified end to end: ping lines appear at exactly 20-second intervals on an
idle session, pong lines remain absent.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 experimental/ssh/internal/proxy/client.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go
index 47c4cb09b6d..a71d53c692b 100644
--- a/experimental/ssh/internal/proxy/client.go
+++ b/experimental/ssh/internal/proxy/client.go
@@ -106,6 +106,10 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque
 						// error returned here would cancel the session it exists to preserve.
 						// The receiving loop notices a genuinely dead connection within one read.
 						log.Debugf(gCtx, "Failed to send websocket keepalive ping: %v", err)
+					} else {
+						// The driver proxy does not return pongs (verified end to end), so this
+						// line is the only evidence in a customer's log that pings were flowing.
+						log.Debugf(gCtx, "Sent websocket keepalive ping")
 					}
 				}
 			}

From bede5759e2b088364391887a98fdb10445507e44 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 24 Aug 2026 14:56:08 +0000
Subject: [PATCH 3/3] Send keepalive pings with WriteControl, off the handover
 mutex
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review of #6358 pointed out that routing pings through sendMessage puts an
unbounded write on the connection's shared write path. WriteMessage takes the
handover mutex and sets no deadline, and close() needs that same mutex, so a
ping that parks on a stalled or half-open socket — a full send buffer, no RST —
holds up the closing handshake until the kernel abandons its retransmits,
roughly 15 minutes with Linux defaults. Context cancellation cannot interrupt a
goroutine parked in a blocking write.

The probability is low for a purely idle session, whose 8-byte control frames
go into a near-empty buffer, but not for the path this feature exists to serve:
data flows, the peer vanishes mid-transfer leaving unacked bytes in the buffer,
the session goes idle, and the next ping parks. That is the feature's own
domain, so the exposure belongs to this change even though the hazard predates
it on the data path.

Pings now go out with WriteControl on the loaded connection, taking no handover
mutex. gorilla explicitly permits WriteControl concurrently with the data
writes, and its deadline bounds both the wait for the connection's write lock
and the socket write itself, so a stalled ping can hold that lock for at most
proxyPingWriteTimeout instead of minutes. The handover path is fully decoupled:
a ping that ticks during a rotation goes to the connection being replaced and
may simply fail, which is already non-fatal.

This drops the "single concurrent writer" rationale for the mutex, which
applies to WriteMessage and not to WriteControl, and adds a write deadline the
original design ruled out. The rule it was protecting — a keepalive must never
end a session — is untouched: a timed-out ping is logged at debug and the
ticker continues.

Tests: the two that asserted the mutex path were reworked, since one drove
sendMessage directly and the other's past write deadline is now overridden by
WriteControl's own. A ping is now asserted to complete during an in-flight
handover rather than to block on it, its failure is induced at the socket, and a
new case parks a ping in the socket write and requires the closing handshake to
finish within the ping's deadline.

Verified end to end: 10 pings at exact 20-second intervals across a 200-second
idle session on dogfood, no failures, session intact.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 experimental/ssh/internal/proxy/client.go     |   9 +-
 .../ssh/internal/proxy/keepalive_test.go      | 175 ++++++++++++++----
 experimental/ssh/internal/proxy/proxy.go      |  13 ++
 3 files changed, 158 insertions(+), 39 deletions(-)

diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go
index a71d53c692b..59e13baf979 100644
--- a/experimental/ssh/internal/proxy/client.go
+++ b/experimental/ssh/internal/proxy/client.go
@@ -97,11 +97,10 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque
 				case <-gCtx.Done():
 					return gCtx.Err()
 				case <-ticker.C:
-					// Pings take the same serialised write path as data — gorilla forbids
-					// concurrent writers — so a ping that ticks during a handover blocks until it
-					// finishes and then goes out late. Harmless: a handover establishes a fresh
-					// connection, which resets the peer's idle clock anyway.
-					if err := proxy.sendMessage(websocket.PingMessage, nil); err != nil {
+					// A ping that ticks during a handover goes to the connection being replaced and
+					// may simply fail. Harmless: a handover establishes a fresh connection, which
+					// resets the peer's idle clock anyway, and the next tick uses the new one.
+					if err := proxy.sendPing(); err != nil {
 						// Never fatal. A failed ping knows nothing the data loops don't, and an
 						// error returned here would cancel the session it exists to preserve.
 						// The receiving loop notices a genuinely dead connection within one read.
diff --git a/experimental/ssh/internal/proxy/keepalive_test.go b/experimental/ssh/internal/proxy/keepalive_test.go
index d1d779d4bff..a7fc0c1bf10 100644
--- a/experimental/ssh/internal/proxy/keepalive_test.go
+++ b/experimental/ssh/internal/proxy/keepalive_test.go
@@ -2,10 +2,14 @@ package proxy
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -15,6 +19,55 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// Write modes for pausableConn.
+const (
+	connWriteOK = iota
+	connWriteFail
+	connWritePark
+)
+
+var errTestWriteFailed = errors.New("test: socket write failed")
+
+// pausableConn emulates the socket conditions a keepalive meets on a stalled peer: writes can be
+// made to fail outright, or to park the way a full send buffer does — blocking until the write
+// deadline expires. Wrapping the socket rather than the websocket keeps the production write path
+// (gorilla's own locking and deadline handling) in the test.
+type pausableConn struct {
+	net.Conn
+	mode         atomic.Int32
+	deadline     atomic.Pointer[time.Time]
+	parked       chan struct{}
+	signalParked func()
+}
+
+func newPausableConn(conn net.Conn) *pausableConn {
+	parked := make(chan struct{})
+	return &pausableConn{
+		Conn:         conn,
+		parked:       parked,
+		signalParked: sync.OnceFunc(func() { close(parked) }),
+	}
+}
+
+func (c *pausableConn) SetWriteDeadline(t time.Time) error {
+	c.deadline.Store(&t)
+	return c.Conn.SetWriteDeadline(t)
+}
+
+func (c *pausableConn) Write(p []byte) (int, error) {
+	switch c.mode.Load() {
+	case connWriteFail:
+		return 0, errTestWriteFailed
+	case connWritePark:
+		c.signalParked()
+		if d := c.deadline.Load(); d != nil && !d.IsZero() {
+			time.Sleep(time.Until(*d))
+		}
+		return 0, os.ErrDeadlineExceeded
+	}
+	return c.Conn.Write(p)
+}
+
 // startKeepaliveTestServer stands up a websocket peer that behaves like the SSH server side of the
 // tunnel for an idle session: it sends the first bytes (which the client waits for before it
 // considers the session established), then only reads. Pings it receives are reported on the
@@ -48,17 +101,26 @@ func startKeepaliveTestServer(t *testing.T) (*httptest.Server, <-chan struct{})
 	return server, pings
 }
 
-func keepaliveTestDialer(serverURL string, onConn func(*websocket.Conn)) createWebsocketConnectionFunc {
+// keepaliveTestDialer returns a connection factory for RunClientProxy. onNetConn, when set, receives
+// each connection's underlying socket so a test can control how its writes behave.
+func keepaliveTestDialer(serverURL string, onNetConn func(*pausableConn)) createWebsocketConnectionFunc {
 	wsURL := "ws" + serverURL[4:]
+	dialer := websocket.Dialer{
+		NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+			conn, err := net.Dial(network, addr)
+			if err != nil {
+				return nil, err
+			}
+			wrapped := newPausableConn(conn)
+			if onNetConn != nil {
+				onNetConn(wrapped)
+			}
+			return wrapped, nil
+		},
+	}
 	return func(ctx context.Context, connID string) (*websocket.Conn, error) {
-		conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose
-		if err != nil {
-			return nil, err
-		}
-		if onConn != nil {
-			onConn(conn)
-		}
-		return conn, nil
+		conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose
+		return conn, err
 	}
 }
 
@@ -94,19 +156,19 @@ func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) {
 	server, _ := startKeepaliveTestServer(t)
 	defer server.Close()
 
-	failWrites := func(conn *websocket.Conn) {
-		// A deadline in the past fails every write on this connection, so every ping fails.
-		// Set before the connection is handed to the proxy, so no writer can be in flight.
-		// Reads are unaffected: the connection is otherwise healthy and the session must survive.
-		conn.SetWriteDeadline(time.Now().Add(-time.Second)) // nolint:errcheck
-	}
-
+	var socket atomic.Pointer[pausableConn]
 	src, _ := io.Pipe()
 	done := make(chan error, 1)
 	go func() {
-		done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, failWrites))
+		done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond,
+			keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) }))
 	}()
 
+	// Fail every write once the connection is up, leaving reads healthy: the connection is otherwise
+	// fine and the session must survive the pings that then fail.
+	require.Eventually(t, func() bool { return socket.Load() != nil }, 10*time.Second, 10*time.Millisecond)
+	socket.Load().mode.Store(connWriteFail)
+
 	// Long enough for many pings to be attempted and fail.
 	select {
 	case err := <-done:
@@ -115,18 +177,63 @@ func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) {
 	}
 }
 
-// TestKeepalivePingBlockedByHandoverDoesNotDeadlock asserts the invariant the keepalive design
-// rests on: a ping sent through the proxy's serialised write path blocks for the duration of a
-// handover, and a handover waits on the receiving loop rather than that write path, so the two
-// cannot deadlock each other.
-func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) {
+// TestKeepalivePingParkedInWriteDoesNotStallClose covers the hazard of writing from a third
+// goroutine: a ping on a stalled or half-open connection parks in the socket write while holding the
+// websocket's write lock, which the closing handshake also needs. Its deadline is what keeps that
+// from lasting until the kernel abandons its retransmits, minutes later.
+//
+// Scoped to the ping's own contribution. pausableConn parks only for as long as the write's deadline,
+// and a write whose caller set none fails at once, so the pre-existing unbounded park on the data
+// path is out of the picture. The session's own shutdown cannot be measured here either: it waits on
+// the receiving loop, which unblocks only when the peer reacts to the close frame, and a peer that
+// has silently gone away never does — with or without a keepalive.
+func TestKeepalivePingParkedInWriteDoesNotStallClose(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
+	server, _ := startKeepaliveTestServer(t)
+	defer server.Close()
+
+	var socket atomic.Pointer[pausableConn]
+	proxy := newProxyConnection(keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) }))
+	require.NoError(t, proxy.connect(ctx))
+	require.NotNil(t, socket.Load())
+
+	socket.Load().mode.Store(connWritePark)
+	pingDone := make(chan error, 1)
+	go func() {
+		pingDone <- proxy.sendPing()
+	}()
+	select {
+	case <-socket.Load().parked:
+	case <-time.After(10 * time.Second):
+		t.Fatal("the keepalive ping did not park in the socket write")
+	}
+
+	start := time.Now()
+	closeErr := proxy.close()
+	require.Less(t, time.Since(start), proxyPingWriteTimeout+5*time.Second,
+		"the closing handshake waited on the parked keepalive ping for longer than its write deadline allows")
+	// The write itself fails, which close() reports; the point is that it was not held indefinitely.
+	require.Error(t, closeErr)
+
+	select {
+	case err := <-pingDone:
+		require.Error(t, err, "a parked ping write must end in an error, not succeed")
+	case <-time.After(10 * time.Second):
+		t.Fatal("the keepalive ping never returned from its parked write")
+	}
+}
+
+// TestKeepalivePingDuringHandoverDoesNotDisruptIt asserts the two periodic behaviours of the tunnel
+// stay independent: a ping sent while a handover is in flight neither waits for the handover nor
+// breaks it. The keepalive writes from a third goroutine, so nothing else guarantees this.
+func TestKeepalivePingDuringHandoverDoesNotDisruptIt(t *testing.T) {
 	ctx := t.Context()
 
 	server := setupTestServer(ctx, t)
 	defer server.Cleanup()
 
-	// Holds the handover open at the point where it has taken the write path but has not yet
-	// completed, which is when a ping must block rather than break the handover.
+	// Holds the handover open at the point where it has taken the handover mutex but has not yet
+	// swapped the connection, which is when a ping must neither block nor interfere.
 	handoverDialing := make(chan struct{})
 	releaseHandover := make(chan struct{})
 	var dials atomic.Int32
@@ -147,13 +254,14 @@ func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) {
 
 	pingDone := make(chan error, 1)
 	go func() {
-		pingDone <- client.Proxy.sendMessage(websocket.PingMessage, nil)
+		pingDone <- client.Proxy.sendPing()
 	}()
 
 	select {
 	case err := <-pingDone:
-		t.Fatalf("ping was written while a handover held the write path: %v", err)
-	case <-time.After(100 * time.Millisecond):
+		require.NoError(t, err)
+	case <-time.After(proxyPingWriteTimeout):
+		t.Fatal("keepalive ping waited on the in-flight handover")
 	}
 
 	close(releaseHandover)
@@ -162,12 +270,11 @@ func TestKeepalivePingBlockedByHandoverDoesNotDeadlock(t *testing.T) {
 	case err := <-handoverDone:
 		require.NoError(t, err)
 	case <-time.After(10 * time.Second):
-		t.Fatal("handover deadlocked while a keepalive ping waited on the write path")
-	}
-	select {
-	case err := <-pingDone:
-		require.NoError(t, err)
-	case <-time.After(10 * time.Second):
-		t.Fatal("keepalive ping never completed after the handover finished")
+		t.Fatal("handover did not complete after a concurrent keepalive ping")
 	}
+
+	// The tunnel still carries data on the connection the handover installed.
+	_, err := client.Input.Write(createTestMessage("client", 1))
+	require.NoError(t, err)
+	require.NoError(t, server.Output.WaitForWrite(createTestMessage("client", 1)))
 }
diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go
index e761cbf54c4..90526e17ea9 100644
--- a/experimental/ssh/internal/proxy/proxy.go
+++ b/experimental/ssh/internal/proxy/proxy.go
@@ -29,6 +29,10 @@ const (
 	proxyHandoverInitTimeout = 30 * time.Second
 	// Timeout for the handover process, when accepted by the server.
 	proxyHandoverAcceptTimeout = 25 * time.Second
+	// Bounds how long a keepalive ping may hold the websocket's write lock. A stalled or half-open
+	// connection parks a write until the kernel gives up retransmitting (~15 minutes with Linux
+	// defaults), and close() and the sending loop need that same lock, so the ping caps its wait.
+	proxyPingWriteTimeout = 5 * time.Second
 )
 
 // handoverCoordination holds the context and channels used to coordinate a single handover operation
@@ -204,6 +208,15 @@ func (pc *proxyConnection) sendMessage(mt int, data []byte) error {
 	return conn.WriteMessage(mt, data)
 }
 
+// sendPing writes a keepalive ping on the current connection. Unlike sendMessage it takes neither
+// the handover mutex nor an unbounded wait: gorilla permits WriteControl concurrently with the data
+// writes, and its deadline bounds how long a stalled socket holds the connection's write lock, which
+// close() and the sending loop also need.
+func (pc *proxyConnection) sendPing() error {
+	conn := pc.conn.Load()
+	return conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(proxyPingWriteTimeout))
+}
+
 func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) error {
 	for {
 		if ctx.Err() != nil {