experimental/ssh: survive a dropped tunnel connection instead of ending the session

.nextchanges/cli/ssh-tunnel-connection-drops.md

@@ -0,0 +1 @@
+* `ssh connect` sessions no longer end when the tunnel's websocket connection is lost. The CLI reattaches to the running session and replays the bytes that were missed, so the shell and everything running in it stay intact, and a transient failure to open a replacement connection for the periodic auth refresh is retried rather than ending the session. Reattaching requires an SSH server started by a CLI that supports it; against an older server the connection behaves as before. ([#6558](https://github.com/databricks/cli/pull/6558))

experimental/ssh/internal/client/client.go

@@ -6,6 +6,7 @@ import (
 	_ "embed"
 	"encoding/base64"
 	"encoding/hex"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -450,7 +451,12 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 	outcome.isSuccess = true
 
 	if opts.ProxyMode {
-		return runSSHProxy(ctx, client, serverPort, clusterID, opts)
+		proxyErr := runSSHProxy(ctx, client, serverPort, clusterID, opts)
+		// isSuccess stays true - the tunnel was established - so the category is what says
+		// whether the session ran to completion or was cut short, and why. Without it a
+		// mid-session drop is indistinguishable from a clean exit in telemetry.
+		outcome.errorCategory = proxySessionEndCategory(proxyErr)
+		return proxyErr
 	} else if opts.IDE != "" {
 		return runIDE(ctx, client, userName, keyPath, serverPort, clusterID, opts)
 	} else {
@@ -898,13 +904,46 @@ func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, use
 }
 
 func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, serverPort int, clusterID string, opts ClientOptions) error {
-	createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) {
-		return createWebsocketConnection(ctx, client, connID, clusterID, serverPort, opts.Liteswap)
+	resumable := serverSupportsResume(ctx, client, clusterID, serverPort, opts.Liteswap)
+	if !resumable {
+		log.Infof(ctx, "The SSH server does not support session resume, a dropped connection will end the session")
+	}
+	createConn := func(ctx context.Context, req proxy.DialRequest) (*websocket.Conn, error) {
+		req.ResumeCapable = resumable
+		return createWebsocketConnection(ctx, client, req, clusterID, serverPort, opts.Liteswap)
 	}
 	requestHandoverTick := func() <-chan time.Time {
 		return time.After(opts.HandoverTimeout)
 	}
-	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, createConn)
+	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, resumable, createConn)
+}
+
+// serverSupportsResume reports whether the running SSH server speaks the resume protocol.
+func serverSupportsResume(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, serverPort int, liteswap string) bool {
+	req, err := newDriverProxyRequest(ctx, client, clusterID, serverPort, "capabilities", liteswap)
+	if err != nil {
+		log.Debugf(ctx, "Failed to build the server capabilities request: %v", err)
+		return false
+	}
+	httpClient := &http.Client{Transport: client.Config.HTTPTransport}
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		log.Debugf(ctx, "Failed to query the server capabilities: %v", err)
+		return false
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		log.Debugf(ctx, "The server does not serve /capabilities (status %d)", resp.StatusCode)
+		return false
+	}
+	var capabilities struct {
+		Resume bool `json:"resume"`
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&capabilities); err != nil {
+		log.Debugf(ctx, "Failed to decode the server capabilities: %v", err)
+		return false
+	}
+	return capabilities.Resume
 }
 
 // accessModeUILabel maps a cluster's access mode to the name shown in the Databricks UI.
@@ -1278,11 +1317,22 @@ func sshExtensionErrorCategory(err error) protos.SshTunnelErrorCategory {
 	return protos.SshTunnelErrorCategoryUnknown
 }
 
-// category returns the error category to report. An interrupted attempt means the user gave
-// up, whichever call happened to observe it first, so it wins over the category recorded at
-// the failure site. An unattributed failure is reported as UNKNOWN so that it stays countable.
+// category returns the error category to report. Once the tunnel is up nothing that follows is a
+// connection failure, so an established session reports only how it ended, and only the proxy can
+// say that. For a connection attempt, an interruption means the user gave up, whichever call
+// happened to observe it first, so it wins over the category recorded at the failure site; an
+// unattributed attempt is reported as UNKNOWN so that it stays countable.
 func (o connectOutcome) category() protos.SshTunnelErrorCategory {
-	if o.isSuccess || o.err == nil {
+	if o.err == nil {
+		return protos.SshTunnelErrorCategoryUnspecified
+	}
+	if o.isSuccess {
+		// proxySessionEndCategory is the only thing that sets a category this late, so an empty
+		// one means the session simply ended: an interruption, an ordinary exit, or a non-zero
+		// exit from the ssh client or the user's own remote command. None is a tunnel failure.
+		if o.errorCategory != "" {
+			return o.errorCategory
+		}
 		return protos.SshTunnelErrorCategoryUnspecified
 	}
 	if errors.Is(o.ctxErr, context.Canceled) || errors.Is(o.err, context.Canceled) {
@@ -1294,6 +1344,27 @@ func (o connectOutcome) category() protos.SshTunnelErrorCategory {
 	return o.errorCategory
 }
 
+// proxySessionEndCategory attributes how a proxy-mode session ended. A dropped websocket is
+// checked first on purpose: a drop that lands during a handover surfaces from either the
+// receiving loop or the handover goroutine, whichever the errgroup records first, and it
+// should be counted as a drop in both cases. HANDOVER_FAILED is then only the handover's own
+// failures. An unrecognised error is left unattributed - normalizeProxyError already maps a
+// clean finish and a user interrupt to nil.
+func proxySessionEndCategory(err error) protos.SshTunnelErrorCategory {
+	switch {
+	case err == nil:
+		return ""
+	case errors.Is(err, proxy.ErrWebsocketDropped):
+		return protos.SshTunnelErrorCategoryWebsocketDropped
+	case errors.Is(err, proxy.ErrHandoverFailed):
+		return protos.SshTunnelErrorCategoryHandoverFailed
+	case errors.Is(err, proxy.ErrConnectFailed):
+		return protos.SshTunnelErrorCategoryWebsocketConnectFailed
+	default:
+		return ""
+	}
+}
+
 func logSshTunnelEvent(ctx context.Context, opts ClientOptions, outcome connectOutcome) {
 	telemetry.Log(ctx, protos.DatabricksCliLog{
 		SshTunnelEvent: buildSshTunnelEvent(opts, outcome),

experimental/ssh/internal/client/client_internal_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/databricks/cli/experimental/ssh/internal/proxy"
 	"github.com/databricks/cli/experimental/ssh/internal/sshconfig"
 	"github.com/databricks/cli/experimental/ssh/internal/vscode"
 	"github.com/databricks/cli/libs/cmdio"
@@ -525,6 +526,17 @@ func TestConnectOutcomeCategory(t *testing.T) {
 			outcome: connectOutcome{isSuccess: true, err: errFailed},
 			want:    protos.SshTunnelErrorCategoryUnspecified,
 		},
+		{
+			// A session end the proxy did attribute must survive isSuccess, or a mid-session
+			// drop is indistinguishable from a clean exit.
+			name: "attributed session end after a successful connection keeps its category",
+			outcome: connectOutcome{
+				isSuccess:     true,
+				errorCategory: protos.SshTunnelErrorCategoryWebsocketDropped,
+				err:           errFailed,
+			},
+			want: protos.SshTunnelErrorCategoryWebsocketDropped,
+		},
 		{
 			name:    "attributed failure keeps its category",
 			outcome: connectOutcome{errorCategory: protos.SshTunnelErrorCategoryIDECommandNotOnPath, err: errFailed},
@@ -582,6 +594,18 @@ func TestConnectOutcomeCategory(t *testing.T) {
 			outcome: connectOutcome{isSuccess: true, ctxErr: context.Canceled, err: errFailed},
 			want:    protos.SshTunnelErrorCategoryUnspecified,
 		},
+		{
+			// ...but a session end the proxy did attribute outranks the interruption, or a drop
+			// that happened to coincide with the user giving up would be lost.
+			name: "an attributed session end wins over an interruption",
+			outcome: connectOutcome{
+				isSuccess:     true,
+				ctxErr:        context.Canceled,
+				errorCategory: protos.SshTunnelErrorCategoryWebsocketDropped,
+				err:           errFailed,
+			},
+			want: protos.SshTunnelErrorCategoryWebsocketDropped,
+		},
 	}
 
 	for _, tt := range tests {
@@ -635,6 +659,53 @@ func TestSshExtensionErrorCategory(t *testing.T) {
 	}
 }
 
+func TestProxySessionEndCategory(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want protos.SshTunnelErrorCategory
+	}{
+		{
+			name: "clean finish is not attributed",
+			err:  nil,
+			want: "",
+		},
+		{
+			name: "dropped websocket",
+			err:  fmt.Errorf("wrapped: %w", proxy.ErrWebsocketDropped),
+			want: protos.SshTunnelErrorCategoryWebsocketDropped,
+		},
+		{
+			name: "handover failure",
+			err:  fmt.Errorf("wrapped: %w", proxy.ErrHandoverFailed),
+			want: protos.SshTunnelErrorCategoryHandoverFailed,
+		},
+		{
+			name: "connect failure",
+			err:  fmt.Errorf("wrapped: %w", proxy.ErrConnectFailed),
+			want: protos.SshTunnelErrorCategoryWebsocketConnectFailed,
+		},
+		{
+			// A drop landing during a handover surfaces from either the receiving loop or the
+			// handover goroutine, so it must be counted as a drop either way.
+			name: "a drop during a handover counts as a drop",
+			err:  errors.Join(proxy.ErrHandoverFailed, proxy.ErrWebsocketDropped),
+			want: protos.SshTunnelErrorCategoryWebsocketDropped,
+		},
+		{
+			name: "an unrecognised error is left unattributed",
+			err:  errors.New("something else"),
+			want: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, proxySessionEndCategory(tt.err))
+		})
+	}
+}
+
 func TestBuildSshTunnelEventReportsErrorCategory(t *testing.T) {
 	got := buildSshTunnelEvent(ClientOptions{ConnectionName: "my-conn", IDE: "vscode"}, connectOutcome{
 		errorCategory: protos.SshTunnelErrorCategoryIDECommandNotOnPath,

experimental/ssh/internal/client/websockets.go

@@ -5,14 +5,16 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strconv"
 
+	"github.com/databricks/cli/experimental/ssh/internal/proxy"
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/gorilla/websocket"
 )
 
-func createWebsocketConnection(ctx context.Context, client *databricks.WorkspaceClient, connID, clusterID string, serverPort int, liteswap string) (*websocket.Conn, error) {
-	proxyURL, err := getProxyURL(ctx, client, connID, clusterID, serverPort)
+func createWebsocketConnection(ctx context.Context, client *databricks.WorkspaceClient, dial proxy.DialRequest, clusterID string, serverPort int, liteswap string) (*websocket.Conn, error) {
+	proxyURL, err := getProxyURL(ctx, client, dial, clusterID, serverPort)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get proxy URL: %w", err)
 	}
@@ -38,19 +40,19 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace
 	return conn, nil
 }
 
-func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, connID, clusterID string, serverPort int) (string, error) {
+func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, dial proxy.DialRequest, clusterID string, serverPort int) (string, error) {
 	workspaceID, err := auth.ResolveWorkspaceID(ctx, client)
 	if err != nil {
 		return "", fmt.Errorf("failed to get current workspace ID: %w", err)
 	}
-	return buildProxyWebsocketURL(client.Config.Host, workspaceID, clusterID, serverPort, connID)
+	return buildProxyWebsocketURL(client.Config.Host, workspaceID, clusterID, serverPort, dial)
 }
 
 // buildProxyWebsocketURL builds the driver-proxy websocket URL for an SSH tunnel.
 //
 // The scheme follows the host (http -> ws, else wss) instead of being hardcoded to
 // wss, so the tunnel is also diallable against the plaintext local test server.
-func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, connID string) (string, error) {
+func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, dial proxy.DialRequest) (string, error) {
 	u, err := url.Parse(host)
 	if err != nil {
 		return "", fmt.Errorf("failed to parse host %q: %w", host, err)
@@ -65,6 +67,15 @@ func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int,
 	// the driver-proxy endpoint and uses an "o" path segment regardless of
 	// whether the workspace ID itself is the legacy or new shape.
 	u.Path = fmt.Sprintf("/driver-proxy-api/o/%s/%s/%d/ssh", workspaceID, clusterID, serverPort)
-	u.RawQuery = url.Values{"id": {connID}}.Encode()
+	query := url.Values{"id": {dial.ConnID}}
+	if dial.ResumeCapable {
+		// Sending "delivered" at all is what tells the server this client speaks the resume
+		// protocol, so it buffers its own output for replay from the start of the session.
+		query.Set("delivered", strconv.FormatInt(dial.Delivered, 10))
+		if dial.Reattach {
+			query.Set("reattach", "1")
+		}
+	}
+	u.RawQuery = query.Encode()
 	return u.String(), nil
 }

experimental/ssh/internal/client/websockets_test.go

@@ -3,6 +3,7 @@ package client
 import (
 	"testing"
 
+	"github.com/databricks/cli/experimental/ssh/internal/proxy"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -11,23 +12,48 @@ func TestBuildProxyWebsocketURL(t *testing.T) {
 	tests := []struct {
 		name string
 		host string
+		dial proxy.DialRequest
 		want string
 	}{
 		{
 			name: "https host is dialed over wss",
 			host: "https://my-workspace.cloud.databricks.test",
+			dial: proxy.DialRequest{ConnID: "conn-1"},
 			want: "wss://my-workspace.cloud.databricks.test/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1",
 		},
 		{
 			name: "plaintext http host is dialed over ws",
 			host: "http://127.0.0.1:8080",
+			dial: proxy.DialRequest{ConnID: "conn-1"},
 			want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1",
 		},
+		{
+			// A server that does not speak the resume protocol must see the URL it always saw,
+			// so no resume parameters leak out when the capability probe said no.
+			name: "a non-resumable dial carries no resume parameters",
+			host: "http://127.0.0.1:8080",
+			dial: proxy.DialRequest{ConnID: "conn-1", Delivered: 4096, Reattach: true},
+			want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1",
+		},
+		{
+			// The offset travels on every dial, not just a reattach: its presence is what tells
+			// the server to start buffering its own output for replay.
+			name: "a resumable dial always carries the delivered offset",
+			host: "http://127.0.0.1:8080",
+			dial: proxy.DialRequest{ConnID: "conn-1", ResumeCapable: true},
+			want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?delivered=0&id=conn-1",
+		},
+		{
+			name: "a reattach states its intent and its offset",
+			host: "http://127.0.0.1:8080",
+			dial: proxy.DialRequest{ConnID: "conn-1", ResumeCapable: true, Delivered: 4096, Reattach: true},
+			want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?delivered=4096&id=conn-1&reattach=1",
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := buildProxyWebsocketURL(tt.host, "900800700600", "1234-567890-abc", 7772, "conn-1")
+			got, err := buildProxyWebsocketURL(tt.host, "900800700600", "1234-567890-abc", 7772, tt.dial)
 			require.NoError(t, err)
 			assert.Equal(t, tt.want, got)
 		})