Fix SSH session resume with bounded backpressure

.nextchanges/cli/ssh-resume-backpressure.md

@@ -0,0 +1 @@
+* Preserve SSH sessions across temporary tunnel disconnects, with bounded replay and backpressure for large transfers. ([#6650](https://github.com/databricks/cli/pull/6650))

experimental/ssh/internal/client/capabilities_test.go

@@ -0,0 +1,46 @@
+package client
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/stretchr/testify/require"
+)
+
+// The capabilities probe runs before the tunnel is dialled, on every proxy-mode connect. A driver
+// proxy that accepts the request and never answers must not hold the connect path open: resume is
+// an optimisation, and failing to learn about it has to cost bounded time, not the whole session.
+func TestServerSupportsResumeGivesUpOnAStalledServer(t *testing.T) {
+	restore := capabilitiesProbeTimeout
+	capabilitiesProbeTimeout = time.Second
+	defer func() { capabilitiesProbeTimeout = restore }()
+
+	stalled := make(chan struct{})
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("GET /driver-proxy-api/o/123/cluster/7772/capabilities", func(w http.ResponseWriter, r *http.Request) {
+		<-stalled
+	})
+	server := httptest.NewServer(mux)
+	// LIFO: release the stalled handler first, so Close does not wait on it.
+	defer server.Close()
+	defer close(stalled)
+
+	client, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "test-token", WorkspaceID: "123", AuthType: "pat"})
+	require.NoError(t, err)
+
+	done := make(chan bool, 1)
+	go func() {
+		done <- serverSupportsResume(t.Context(), client, "cluster", 7772, "")
+	}()
+
+	select {
+	case resumable := <-done:
+		require.False(t, resumable, "a server that never answered cannot be assumed to support resume")
+	case <-time.After(30 * time.Second):
+		t.Fatal("the capabilities probe never returned: it has no timeout of its own, so a stalled driver proxy blocks the connect path indefinitely")
+	}
+}

experimental/ssh/internal/client/client.go

@@ -6,6 +6,7 @@ import (
 	_ "embed"
 	"encoding/base64"
 	"encoding/hex"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -493,7 +494,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, knownHostsPath, serverPort, clusterID, opts)
 	} else {
@@ -973,13 +979,52 @@ 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)
+}
+
+// capabilitiesProbeTimeout caps the pre-connect capabilities probe. A var so tests can shorten it.
+var capabilitiesProbeTimeout = 10 * time.Second
+
+// 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
+	}
+	// Bounded on its own: this probe runs before the tunnel is dialled, and resume is an
+	// optimisation, so a driver proxy that accepts the request and never answers must cost a
+	// bounded wait rather than the whole connect path.
+	httpClient := &http.Client{Transport: client.Config.HTTPTransport, Timeout: capabilitiesProbeTimeout}
+	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 {
+		ResumeVersion int `json:"resume_version"`
+	}
+	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.ResumeVersion == proxy.ResumeProtocolVersion
 }
 
 // accessModeUILabel maps a cluster's access mode to the name shown in the Databricks UI.
@@ -1351,11 +1396,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) {
@@ -1367,6 +1423,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"
@@ -547,6 +548,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},
@@ -604,6 +616,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 {
@@ -657,6 +681,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

@@ -2,17 +2,20 @@ package client
 
 import (
 	"context"
+	"errors"
 	"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)
 	}
@@ -30,27 +33,35 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace
 	}
 
 	// websocket connection manages lifecycle of the response object, no need to close the body
-	conn, _, err := websocket.DefaultDialer.Dial(req.URL.String(), req.Header) // nolint:bodyclose
+	conn, resp, err := websocket.DefaultDialer.DialContext(ctx, req.URL.String(), req.Header)
+	if resp != nil {
+		resp.Body.Close()
+	}
 	if err != nil {
+		// Only the server's explicit session refusals bypass the retry budget.
+		// Other responses, including 408, can come from an intermediary.
+		if dial.Reattach && resp != nil && (resp.StatusCode == http.StatusConflict || resp.StatusCode == http.StatusGone) {
+			return nil, errors.Join(proxy.ErrReattachRejected, fmt.Errorf("reattach failed (HTTP %d): %w", resp.StatusCode, err))
+		}
 		return nil, fmt.Errorf("failed to establish websocket connection: %w", err)
 	}
 
 	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 +76,16 @@ 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 {
+		query.Set(proxy.ResumeVersionParameter, strconv.Itoa(proxy.ResumeProtocolVersion))
+		// 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
 }